[論文レビュー] Progressive Distillation for Fast Sampling of Diffusion Models

Tim Salimans, Jonathan Ho|arXiv (Cornell University)|Feb 1, 2022

Generative Adversarial Networks and Image Synthesis被引用数 190

ひとこと要約

本論文は、拡散モデルのサンプリングを圧縮するためのプログレッシブ蒸留を導入し、ステップ数を逐次半分に削減して、4ステップという少ないステップ数でも高品質なサンプルを実現しつつ、性能を保つ。

ABSTRACT

Diffusion models have recently shown great promise for generative modeling, outperforming GANs on perceptual quality and autoregressive models at density estimation. A remaining downside is their slow sampling time: generating high quality samples takes many hundreds or thousands of model evaluations. Here we make two contributions to help eliminate this downside: First, we present new parameterizations of diffusion models that provide increased stability when using few sampling steps. Second, we present a method to distill a trained deterministic diffusion sampler, using many steps, into a new diffusion model that takes half as many sampling steps. We then keep progressively applying this distillation procedure to our model, halving the number of required sampling steps each time. On standard image generation benchmarks like CIFAR-10, ImageNet, and LSUN, we start out with state-of-the-art samplers taking as many as 8192 steps, and are able to distill down to models taking as few as 4 steps without losing much perceptual quality; achieving, for example, a FID of 3.0 on CIFAR-10 in 4 steps. Finally, we show that the full progressive distillation procedure does not take more time than it takes to train the original model, thus representing an efficient solution for generative modeling using diffusion at both train and test time.

研究の動機と目的

• 条件なしおよびクラス条件付きの画像生成における拡散モデルの遅いサンプリングのボトルネックを動機づけて対処する。
• 遅くて多ステップの教師モデルから速い生徒モデルへ精度を移す蒸留法を提案する。
• 進行的蒸留がステップを減らしつつ、標準ベンチマーク（CIFAR-10、ImageNet、LSUN）でサンプル品質を維持することを実証する。
• 劣化なしの迅速な蒸留を支援する安定な拡散のパラメータ化と損失重み付け戦略を提供する。
• 全体の蒸留プロセスが元のモデルの訓練と比較して計算効率が高いことを示す。

提案手法

• プログレッシブ蒸留を定式化し、1つの生徒ステップで2つの教師DDIMステップに一致するよう生徒モデルを訓練し、その後ステップ数を反復的に半分にする。
• 教師が与えられた z_t から2つのDDIMステップを実行し、それを反転させて生徒に鋭いターゲットを得る蒸留ターゲットを用いる。
• 拡散デノイジングモデルの異なるパラメータ化（direct x、epsilon、または v）と、信号対ノイズ比が低いときにも情報を保つ損失重み付けを用いて実験する。
• コサインノイズスケジュールと BigGAN 風のアップ/ダウンサンプリングを備えた U-Net アーキテクチャを用い、標準的な拡散モデル訓練手順に従って訓練する。
• 蒸留モデルを DDIM および確率的ベースラインと比較し、データセット全体で FID（適用可能な場合は IS）を報告する。
• オープンソース実装と詳細な再現性ノートを提供する。

実験結果

リサーチクエスチョン

• RQ1 プログレッシブ蒸留は、サンプル品質の大幅な低下なしに拡散モデルのサンプリングステップ数を意味のある程度減らせるか？
• RQ2 条件なし・クラス条件付き生成の双方で安定し有効な蒸留を生み出すパラメータ化と損失重み付け戦略はどれか？
• RQ3 標準ベンチマーク（CIFAR-10、ImageNet、LSUN）において、蒸留モデルの品質は既存の高速サンプリング手法とどう比較されるか？
• RQ4 元の拡散モデルを訓練することと比べたプログレッシブ蒸留の計算コストはどのくらいで、実際に効率的か？

主な発見

• 蒸留モデルは CIFAR-10 で 4 サンプリングステップという少ないステップ数でも高品質なサンプルを生成でき、競争力のある FID スコアを達成する。
• プログレッシブ蒸留は各反復でステップ数を半分にすることで大幅なスピードアップを生み出し、全体の手順が元のモデルの訓練時間を超えない。
• 異なる安定なパラメータ化（direct x、epsilon、または v）と損失重み付け（SNRベースの変種）は堅牢な性能を生み出すが、アブレーションで1つ不安定な組み合わせが特定された。
• CIFAR-10、64×64 ImageNet、128×128 LSUN のベンチマークで、蒸留モデルは同程度またはより少ないステップ数で、特に4～8ステップ領域で高速ベースラインを上回るか対等に競う。
• 蒸留プロセスは効率的であり、総時間は元のモデルを訓練するのに要する時間を超えない。このため推定時の効率化が実用的となる。
• 蒸留モデルは確率的サンプラーとも併用でき、DDIMと確率的ベースラインの中間の性能を得る。