Diffusion2GAN: 拡散モデルを条件付きGANに知識蒸留
3つの要点
✔️ 拡散モデルはLAIONなどの難しいデータセットで優れた画像生成品質を示しているが、高品質な結果には数十から数百のサンプリングステップが必要
✔️ 拡散モデルを条件付きGANに知識蒸留という方針で、ワンステップで高品質の生成モデルの手法を提案
✔️ 生成画像の画質を向上するために、新しい知識蒸留の損失と識別機を提案
Distilling Diffusion Models into Conditional GANs
written by Minguk Kang, Richard Zhang, Connelly Barnes, Sylvain Paris, Suha Kwak, Jaesik Park, Eli Shechtman, Jun-Yan Zhu, Taesung Park
(Submitted on 9 May 2024)
Comments: Published on arxiv.
Subjects: Computer Vision and Pattern Recognition (cs.CV)
code：
本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。
はじめに
拡散モデルは、LAIONなどの難しいデータセットで優れた画像生成品質を示していますが、高品質な結果には数十から数百のサンプリングステップが必要です。そのため、DALL·E 2、Imagen、Stable Diffusionなどのモデルは高レイテンシーでリアルタイム対話が難しいです。もし1ステップのモデルがあれば、テキストから画像生成のユーザー体験が向上し、3Dや動画アプリケーションでも利用可能になります。
単純な解決策として、1ステップモデルをゼロからトレーニングする方法があります。GANはシンプルなドメインで優れた1ステップモデルですが、大規模で多様なデータセットでのテキストから画像生成にはまだ課題があります。この課題は、ノイズと自然画像の対応を見つけ、ノイズから画像へのマッピングを最適化する2つのタスクを教師無しで行う必要があるためです。
今回の解説論文のアイデアは、これらのタスクを一つずつ解決することです。まず、事前学習済みの拡散モデルでノイズと画像の対応を見つけ、その後、条件付きGANにペアありの画像変換フレームワークでノイズから画像へのマッピングを行わせます。このアプローチにより、拡散モデルの高品質と条件付きGANの迅速なマッピングの利点を両立させることができます。
実験では、提案されたDiffusion2GANフレームワークを使用して、Stable Diffusion 1.5を単一ステップの条件付きGANモデルに蒸留しました。Diffusion2GANは、他の蒸留方法よりもターゲットの拡散モデルに固有のノイズから画像への対応をより良く学習できます。さらに、UFOGenやDMDを凌駕し、COCO2014ベンチマークで優れた結果を示しました。特に、SDXLの蒸留ではSDXL-TurboやSDXL-Lightningを上回りました。
提案手法
全体の概要
提案手法全体の概要は図1に示されています。最初に、拡散モデルの出力潜在変数と入力ノイズ、プロンプトを収集します。次に、ジェネレータは、E-LatentLPIPS損失とGAN損失を使用して、ノイズとプロンプトを目標の潜在変数にマッピングするようにトレーニングされます。ジェネレータの出力はRGBピクセルにデコードできますが、これはトレーニング中に実行されない高計算量の操作です。E-LatentLPIPS損失とGAN損失についはは次の説で詳しく説明します。
E-LatentLPIPS損失
従来の知識蒸留の損失は次の式のようになります。この損失をそのまま利用可能（図2−b）ですが、ピクセル空間のために設計されたので、Decoderで潜在空間からピクセル空間に変更する必要です。
これは計算量が多い操作であるため、ピクセルへのデコードが必要なく、潜在空間での知覚的距離を直接計算する方法が必要です。
これを実現するために、Zhangらの手順に従い、Stable Diffusionの潜在空間でImageNet上でVGGネットワークをトレーニングしました。潜在空間はすでに8倍にダウンサンプリングされているため、アーキテクチャをわずかに変更しました。その後、BAPPSデータセットを使用して中間特徴を線形にキャリブレーションし、潜在空間で動作する関数を得ました。
$d_{LatentLPIPS}(x_0, x_1) = ℓ(F(x_0), F(x_1))$
しかし、図2−cに示したように、LatentLPIPSを蒸留の新しい損失関数として直接適用すると、波状のパッチ状のアーティファクトが生じることを確認できます。
E-LPIPSからのインスピレーションを受けて、生成された潜在変数と目標潜在変数の両方にランダムな微分可能な拡張、一般的な幾何学的変換、およびカットアウトを適用します。各エポックで、ランダムな拡張が生成された潜在変数と目標潜在変数の両方に適用されます。単一画像最適化に適用された場合、アンサンブル戦略はほぼ完璧に目標画像を再構築します（図2-d）を参照）。新しい損失関数は、Ensembled-LatentLPIPS、またはE-LatentLPIPSと短縮されます。Tはランダムにサンプリングされた拡張です。
条件付き拡散識別器
前節では、拡散蒸留を、ペアのノイズから潜在変数への変換タスクとしてフレーム化することが可能であることを明らかにしました。ペア画像から画像への変換のための条件付きGANの効果に着想を得て、条件付きの識別器を使用します。この識別器の条件には、テキスト記述$c$だけでなく、ジェネレータに提供されるガウスノイズ$z$も含まれます。新しい識別器は、事前学習された拡散の重みを活用しながら、前述の条件を組み込んでいます。具体的には、以下の最小最大目的関数をジェネレータ$G$と識別器$D$に対して最適化します。 識別器の概要は図3に示されています。
実験
蒸留された拡散モデルとの比較
Diffusion2GANと最先端の拡散蒸留モデルを比較し、COCO2014とCOCO2017の結果を表1と表2に示します。InstaFlow-0.9Bは、COCO2014でFID 13.10、COCO2017でFID 23.4を達成し、Diffusion2GANはFID 9.29とFID 19.5を達成します。他の拡散モデルも使用されますが、Diffusion2GANは元のモデルの軌跡に密接に従うように訓練され、高い視覚品質を維持しつつ多様性の崩壊問題を緩和します。ADD-MはViT-g-14テキストエンコーダーを使用し、高いCLIP-5kスコアを達成しますが、Diffusion2GANはこのエンコーダーを使用せずに訓練されています。
視覚的分析
図4では、提案手法をStable Diffusion 1.5、LCM-LoRA、InstaFlowと比較して視覚的に示します。拡散モデルは、クラシファイアフリーガイダンス（CFG）の拡大に伴い、より写真のような画像を生成する傾向があるため、提案手法はDiffusion2GANをSD-CFG-8データセットを使用してトレーニングし、同じガイダンススケール8を使用してStable Diffusion 1.5と比較します。LCM-LoRAとInstaFlowについては、公平な比較を確保するために、それぞれの最良の設定に従います。提案手法の結果は、他の蒸留ベースラインと比較して、画像をよりリアルに生成し、同時にStable Diffusionの教師によって生成されたターゲット画像の全体的なレイアウトを維持することを示しています。
トレーニングのスピード
ODEデータセットの準備コストを含めても、Diffusion2GANは既存の蒸留方法よりも効率的に収束します。CIFAR10データセットでは、トータルトレーニング中の生成器ネットワークの総関数評価数を比較します。提案手法は、LPIPS損失でのトレーニングが、500kの教師出力でConsistency DistillationのFIDを既に超えていることを確認できます（表3）。テキストから画像への合成では、完全バージョンのDiffusion2GANが、InstaFlowと比較して優れたFIDを達成し、かつはるかに少ないGPU日数で実現できています（表4）。
まとめ
今回、事前に訓練されたマルチステップ拡散モデルを、条件付きGANと知覚損失で訓練された単一ステップのジェネレータに蒸留する新しいフレームワークDiffusion2GANを提案しました。提案手法は、生成モデリングを2つのタスクに分割すること（対応関係の特定とマッピングの学習）が、異なる生成モデルを使用してパフォーマンスと計算効率を向上させることを示しています。このシンプルなアプローチは、対話型画像生成を向上させるだけでなく、効率的なビデオや3Dアプリケーションも可能だと考えられます。
しかし、提案手法は、いくつかの制限があります。1つ目は、固定されたガイダンススケールを使っているため、推論時に異なる値を使えないことです。2つ目は、教師モデルの品質に依存するため、性能に限界があります。３目は、学生と教師モデルの規模を拡大すると、依然として多様性の低下が起こることです。今後の研究では、この３つの課題を解決できれば、より実用性が高いモデルとなります。
この記事に関するカテゴリー
