[論文レビュー] On Distillation of Guided Diffusion Models

Chenlin Meng, Rombach, Robin+5 more|arXiv (Cornell University)|Oct 6, 2022

Advanced Neuroimaging Techniques and Applications被引用数 8

ひとこと要約

本稿では、分類器フリーのガイド付き拡散モデルの高速化を目的とした2段階の蒸留フレームワークを提案する。これにより、1〜4ステップのノイズ除去ステップで高忠実度の画像生成が可能になる。まず、条件付きおよび非条件付きの教師モデルの出力を統合して1つの学生モデルに蒸留し、その後段階的にステップ数を減らしてさらに蒸留することで、標準の拡散モデルと比較して最大256倍の高速化を達成した。ImageNetおよびLAIONベンチマークにおいても、FIDおよびISスコアが競争力を持つ結果を得た。

ABSTRACT

Classifier-free guided diffusion models have recently been shown to be highly effective at high-resolution image generation, and they have been widely used in large-scale diffusion frameworks including DALLE-2, Stable Diffusion and Imagen. However, a downside of classifier-free guided diffusion models is that they are computationally expensive at inference time since they require evaluating two diffusion models, a class-conditional model and an unconditional model, tens to hundreds of times. To deal with this limitation, we propose an approach to distilling classifier-free guided diffusion models into models that are fast to sample from: Given a pre-trained classifier-free guided model, we first learn a single model to match the output of the combined conditional and unconditional models, and then we progressively distill that model to a diffusion model that requires much fewer sampling steps. For standard diffusion models trained on the pixel-space, our approach is able to generate images visually comparable to that of the original model using as few as 4 sampling steps on ImageNet 64x64 and CIFAR-10, achieving FID/IS scores comparable to that of the original model while being up to 256 times faster to sample from. For diffusion models trained on the latent-space (e.g., Stable Diffusion), our approach is able to generate high-fidelity images using as few as 1 to 4 denoising steps, accelerating inference by at least 10-fold compared to existing methods on ImageNet 256x256 and LAION datasets. We further demonstrate the effectiveness of our approach on text-guided image editing and inpainting, where our distilled model is able to generate high-quality results using as few as 2-4 denoising steps.

研究の動機と目的

• 分類器フリーのガイド付き拡散モデルが、推論時に2つの別々のモデルを数十〜数百回評価する必要があるという高い計算コストを解消すること。
• 画像品質やガイドランスの柔軟性を損なわずに、蒸留モデルからの高速なサンプリングを可能にする蒸留手法の開発。
• ピクセル空間および潜在変数空間の両方の拡散モデル、特にS待定Diffusionなどの最先端モデルに対しても蒸留技術を拡張すること。
• 最小限のサンプリングステップで効率的なテキストから画像への変換、画像編集、インpainting（穴埋め）を可能にすること。

提案手法

• 教師モデルの条件付きおよび非条件付き拡散モデルの出力の組み合わせを学習する1つの学生モデルを導入する。
• 2段階の蒸留プロセスを採用：まず学生モデルを教師モデルの出力に合わせて学習し、その後段階的にステップ数を減らしてさらに蒸留する。
• 蒸留モデルを用いた推論に、既存の決定的プロセスに加え、提案された確率的サンプリングプロセスを活用する。
• ピクセル空間モデル（例：CIFAR-10およびImageNet 64×64）および潜在変数空間モデル（例：ImageNet 256×256およびLAION 512×512におけるStable Diffusion）の両方の応用を実施する。
• さまざまなガイドランス重みをサポートするように蒸留モデルを訓練することで、多様性と品質のトレードオフを制御可能にする。
• 異なるガイドランス強度において、学生モデルの予測が教師モデルの出力と一致するように、重み付きMSE目的関数を用いて蒸留を行う。

実験結果

リサーチクエスチョン

• RQ1 分類器フリーのガイド付き拡散モデル（2つのモデルからの推論を要する）に対し、蒸留を効果的に適用できるか？
• RQ2 1〜4ステップのノイズ除去ステップで、ガイドランスの柔軟性を保ちながら高品質な画像生成が可能な1つの蒸留モデルを実現できるか？
• RQ3 提案された2段階の蒸留アプローチは、元のモデルと比較してFIDおよびISスコアを維持または向上させられるか？
• RQ4 この蒸留フレームワークは、ピクセル空間および潜在変数空間の両方の拡散モデルアーキテクチャに一般化可能か？
• RQ5 テキストから画像への変換、画像編集、インpainting（穴埋め）といった下流タスクに対しても、最小限のサンプリングステップで効果的に適用可能か？

主な発見

• 4ステップのサンプリングで、ImageNet 64×64およびCIFAR-10において教師モデルと同等のFIDスコアを達成し、推論速度が最大256倍向上した。
• Stable Diffusionなどの潜在変数空間モデルにおいて、サンプリングステップを1〜4に削減し、ImageNet 256×256およびLAION 512×512において、既存手法と比較して少なくとも10倍の高速化を達成した。
• さまざまなガイドランス重み（w = 0 から w = 4）においても、高い視覚的品質を維持しており、多様性と忠実度の間の有効なトレードオフが可能である。
• テキストガイド付きの画像編集およびインpaintingタスクにおいて、2〜4ステップのノイズ除去で高品質な結果を生成し、編集タスクにおける優れた性能を示した。
• 特にステップ数が少ない状況において、決定的サンプリングに比べて確率的サンプリングプロセスがさらに高いサンプル品質をもたらした。
• 本手法は、ピクセル空間および潜在変数空間の両方の分類器フリーのガイド付き拡散モデルに対して、効果的な蒸留を初めて実証した。ベンチマーク全体にわたり一貫した性能向上を達成した。