[論文レビュー] BOOT: Data-free Distillation of Denoising Diffusion Models with Bootstrapping

Jiatao Gu, Shuangfei Zhai+3 more|arXiv (Cornell University)|Jun 8, 2023

Advanced Neuroimaging Techniques and Applications被引用数 4

ひとこと要約

BOOTは、連続するサンプリングステップからのブートストラップを用いて、実際のデータや大量のオフラインデータ生成を必要とせずに、1ステップの学生モデルを訓練する、データフリーな知識蒸留手法を提案する。この手法は、Stable Diffusion や DeepFloyd IF といった大規模なテキスト-to-画像モデルですら、実際のデータなしで教師モデルと同等の画像品質を達成しながら、推論速度を1000倍以上に高速化できる。

ABSTRACT

Diffusion models have demonstrated excellent potential for generating diverse images. However, their performance often suffers from slow generation due to iterative denoising. Knowledge distillation has been recently proposed as a remedy that can reduce the number of inference steps to one or a few without significant quality degradation. However, existing distillation methods either require significant amounts of offline computation for generating synthetic training data from the teacher model or need to perform expensive online learning with the help of real data. In this work, we present a novel technique called BOOT, that overcomes these limitations with an efficient data-free distillation algorithm. The core idea is to learn a time-conditioned model that predicts the output of a pre-trained diffusion model teacher given any time step. Such a model can be efficiently trained based on bootstrapping from two consecutive sampled steps. Furthermore, our method can be easily adapted to large-scale text-to-image diffusion models, which are challenging for conventional methods given the fact that the training sets are often large and difficult to access. We demonstrate the effectiveness of our approach on several benchmark datasets in the DDIM setting, achieving comparable generation quality while being orders of magnitude faster than the diffusion teacher. The text-to-image results show that the proposed approach is able to handle highly complex distributions, shedding light on more efficient generative modeling.

研究の動機と目的

• 拡散モデルの推論速度が遅いため、多数のステップにわたり繰り返しノイズ除去を行うという問題に対処すること。
• 大規模な合成データ生成や実際のトレーニングデータへのアクセスを必要とする従来の蒸留手法の制限を克服すること。
• プライバシー制約や規模の制約によりトレーニングデータにアクセスできない大規模なテキスト-to-画像拡散モデルの効率的蒸留を可能にすること。
• 連続した時刻予測からのブートストラップを活用する、データフリーで効率的なトレーニングフレームワークの開発。
• 完全な拡散教師モデルと同等の品質を実現する高精細な1ステップ生成を達成するとともに、推論時間を著しく短縮すること。

提案手法

• 本手法は、1つのノイズ入力を与えられた場合に、あらかじめ訓練済みの拡散教師モデルの全時刻における出力を予測する時刻条件付きの学生モデルを訓練する。
• PF-ODEから導出された新規の信号-ODE（Signal-ODE）を用い、時刻間での訓練の安定性と一般化性能を向上させる。
• 同じ拡散軌道に沿った決定的マッピングを活用することで、$\bm{x}_{t'}$（$t' > t$）の知識に基づいて$\bm{x}_t$を予測するブートストラップを実現する。
• 推論品質と多様性を向上させるために、定期的に境界損失（boundary loss）を適用するが、頻度を低く抑えることで効率性を維持する。
• 分類器フリー・ガイドランス（CFG）と互換性があるため、条件付きおよびテキスト-to-画像モデルの効果的蒸留が可能である。
• 学生モデルは、実際のデータを一切使用せず、教師モデルの前方伝搬のみを用いてエンドツーエンドで訓練されるため、プライバシーが重要な環境や大規模モデルに適している。

実験結果

リサーチクエスチョン

• RQ1 実際のトレーニングデータや大規模な合成データ生成にアクセスできない状況下でも、1ステップの学生モデルを効果的に訓練できるか？
• RQ2 連続する時刻予測からのブートストラップは、データフリーな環境下での蒸留効率と性能向上にどのように寄与するか？
• RQ3 提案手法は、完全な拡散教師モデルと同等の高精細な画像生成を達成できるか？ また、推論速度を1000倍以上に高速化できるか？
• RQ4 Stable Diffusion や DeepFloyd IF といった複雑で高容量なテキスト-to-画像拡散モデルに対しても、本手法は一般化性を示せるか？
• RQ5 信号-ODEと境界損失は、訓練の安定化とサンプリングの多様性向上にどのような役割を果たすか？

主な発見

• BOOTは、FFHQ、LSUN、ImageNet などの非条件およびクラス条件付き画像生成ベンチマークにおいて、完全な拡散教師モデルと同等の生成品質を達成した。
• 本手法により、教師モデル比で50～1000倍の高速化が実現され、1ステップ推論が可能になった。
• テキスト-to-画像生成において、DeepFloyd IF や Stable Diffusion モデルの両方を、トレーニングデータにアクセスせずに効果的に蒸留し、多様なプロンプトから高品質な画像を生成できた。
• 蒸留済みモデルは、教師モデルと同一の初期ノイズを用いた定性的な比較でも、高い多様性と忠実度を維持していた。
• 4ステップごとに境界損失を適用することで、推論品質と多様性が向上したが、トレーニングコストは著しく増加しなかった。
• 本手法は計算的にも効率的であり、最大10億パラメータのモデルについても、8台のA100 GPUと3～7日間のトレーニングで実現可能で、テキストプロンプト以外のデータロードオーバーヘッドは一切ない。