埋め込みロスによる効率的な拡散蒸留

要旨: 高価な拡散モデルを効率的な少数ステップ生成器へ蒸留する最近の進展は、大きな可能性を示している。しかし、これらの手法は一般に相当な計算資源と長い学習期間を必要とするため、計算資源に制約のある研究者の利用を妨げており、既存の補助損失関数にも顕著な限界がある。回帰損失は学習の前に大規模なデータセットを事前生成する必要があり、学習される学生モデルの性能が教師に制限される。一方でGANベースの損失は学習の不安定性を伴い、慎重なチューニングを要する。本論文では、既存の拡散蒸留手法を補完して生成品質を高め、より小さなバッチサイズで学習を加速する新しい補助損失関数である Embedding Loss（EL）を提案する。ランダムに初期化された多様なネットワーク群から得られる特徴埋め込みを活用することで、ELは蒸留された少数ステップ生成器と元データの間で特徴分布を効果的に整合させる。埋め込み特徴空間上で Maximum Mean Discrepancy（MMD）を計算することにより、ELは頑健な分布整合を保証し、蒸留の過程でサンプルの忠実性と多様性を維持する。分布整合型の蒸留フレームワークの中で、ELは1ステップ生成器に対して強い実証的性能を示す。CIFAR-10データセットにおいて、提案手法は無条件生成で最先端のFID値1.475、条件付き生成で1.380を達成する。CIFAR-10を超えて、さらにELをImageNet、AFHQ-v2、FFHQの各データセットを含む複数のベンチマークおよび蒸留手法にわたって検証し、DMD、DI、CMの蒸留フレームワークを用いることで、既存の1ステップ蒸留手法に対して一貫した改善が得られることを示す。加えて、本手法は学習反復回数を最大80%削減し、計算資源に制約のある環境において拡散ベースの生成モデルを導入するための、より実用的でスケーラブルな解決策を提供する。