HN 日本語サマリー

← 一覧へ戻る
AI・機械学習

拡散蒸留のパラドックス (2024)

The paradox of diffusion distillation (2024) (sander.ai)

6 pointsby peter_d_sherman1 コメント

要約

拡散モデルは、高次元分布からのデータ生成という困難なタスクを、より簡単な多数のノイズ除去タスクに分割することで解決します。しかし、近年、サンプリングステップ数を削減する研究が多く、これは拡散モデルの利点と矛盾するように見えます。この記事では、特に蒸留（教師モデルの予測を用いて学生モデルを訓練する手法）に焦点を当て、拡散モデルのサンプリングステップ数を減らす様々な方法とそのパラドックスについて深く掘り下げます。

全文翻訳

拡散蒸留のパラドックス

拡散モデルは、高次元分布からのデータ生成という困難なタスクを、それぞれがはるかに容易な多くのノイズ除去タスクに分割します。私たちは、モデルに一度に1つのタスクだけを解くように訓練します。サンプリング時には、多くの予測を逐次的に行います。この反復的な洗練こそが、その強力さの源泉です。
…あるいは、そうではないのでしょうか？ 最近の拡散モデルに関する多くの論文は、必要なサンプリングステップ数を削減することに焦点を当てており、単一ステップでのサンプリングを可能にすることを目指す研究さえあります。これは、そもそもこれらのモデルがうまく機能する理由とされる、多くの簡単なステップに分割するという考え方とは逆のように思えます！

この記事では、拡散モデルから良好な結果を得るために必要なサンプリングステップ数を削減できる様々な方法を詳しく見ていきます。特に、蒸留の様々な形態に焦点を当てます。これは、別のモデル（教師）の予測によって監視しながら、新しいモデル（生徒）を訓練する手法です。拡散モデルに対する様々な蒸留手法は、非常に説得力のある結果を生み出しています。

書き始めたときは比較的ハイレベルな内容にするつもりでしたが、拡散モデルの蒸留は少しニッチなテーマであるため、いくつかのことを詳細に説明せざるを得ず、深掘りした内容になりました。以下は目次です。投稿の特定のセクションに直接ジャンプするにはクリックしてください。

  1. 拡散サンプリング：慎重に進め！
  2. 目的を持って入力空間を移動する
  3. 拡散蒸留
  4. しかし、「フリーランチ」はないのか？
  5. 本当に教師が必要か？
  6. データとノイズの間の迷宮をたどる
  7. 締めくくりの考察
  8. 謝辞
  9. 参考文献

拡散サンプリング：慎重に進め！

［画像］

まず最初に、なぜ拡散モデルで良好な結果を得るのに多くのステップが必要なのでしょうか？ 様々な手法が、出力の品質を損なうことなく（あるいは、少なくともそれほど損なうことなく）このステップ数を削減できる理由を理解するには、これについてのより深い理解を深める価値があります。

拡散モデルにおけるサンプリングステップは、以下から構成されます。

  • ノイズを除去するために入力空間のどの方向に移動すべきか、あるいは等価的に、データ分布の下で入力をより尤もらしくするために、その方向を予測する。
  • その方向に小さなステップを取る。

サンプリングアルゴリズムによっては、少量のノイズを追加したり、更新方向を計算するために、より高度なメカニズムを使用したりする場合があります。

私たちは小さなステップしか取らないのは、この予測された方向は局所的にしか意味をなさないからです。それは、尤度が高い入力空間の領域を指していますが、特定のデータポイントを直接指しているわけではありません。したがって、大きなステップを取ると、その高尤度領域の重心に到達してしまい、必ずしもデータ分布の代表的なサンプルとは限りません。これを大まかな推定と考えてください。もしこれが直感に反するように感じるなら、あなたは一人ではありません！高次元空間における確率分布は、しばしば直感に反して振る舞います。これは私が過去に詳細なブログ記事で書いたことです。

具体的には、画像ドメインでは、大きなステップを予測された方向に取ると、入力に多くのノイズが含まれている場合、ぼやけた画像が得られる傾向があります。これは、基本的に多くのあり得る画像の平均に対応するためです。（議論のために、サンプリングアルゴリズムの一部として追加される可能性のあるノイズは意図的に無視しています。）

別の見方をすれば、ノイズは高周波情報（シャープな特徴や細かいディテールに対応します。これも以前に書いたことがあります）を不明瞭にします。この高周波情報に関する不確実性は、すべての可能性がブレンドされた予測をもたらし、結果として高周波情報が完全に失われます。

予測された方向の局所的な妥当性は、微小なステップしか取れないことを意味し、その後、新しい方向を決定するためにモデルを再評価する必要があります。もちろん、これは現実的ではないため、有限ですが小さなステップを取ります。これは、パラメータ空間における機械学習モデルの勾配ベース最適化の働き方に非常に似ていますが、ここでは入力空間で操作しています。モデルトレーニングの場合と同様に、ステップが大きすぎると、最終結果の品質が低下します。

以下は、2次元で入力空間を表す図です。(\mathbf{x}_t) は時間ステップ (t) におけるノイズのある入力を表し、ここではデータ分布から描画されたクリーンな画像 (\mathbf{x}_0) にノイズを追加して構築しました。また、入力をより尤もらしくするために移動すべき方向（拡散モデルによって予測される）も示されています。これは、高尤度領域の重心である (\hat{\mathbf{x}}_0) を指しており、ピンク色で shaded されています。

［画像: 入力空間における高尤度領域と、その領域の重心を指す拡散モデルによって予測された方向を示す図。］

図：入力空間における高尤度領域と、その領域の重心を指す拡散モデルによって予測された方向を示す図。

（非常に高次元の空間を2Dで表現することに関する注意点については、以前のブログ記事「拡散ガイダンスの幾何学」の最初のセクションを参照してください！）

この方向にステップを進め、少量のノイズを追加すると（例えば、DDPM1サンプリングアルゴリズムで行うように）、(\mathbf{x}_{t-1}) に到達します。これは、わずかにノイズが少ない入力画像に対応します。前のサンプリングステップによって不確実性が解消されたため、予測された方向は、より小さく「より具体的な」高尤度領域を指すようになります。これは以下の図に示されています。

［画像: 単一サンプリングステップ後の拡散モデルによって予測された更新された方向と、それが指す対応する高尤度領域を示す図。］

図：単一サンプリングステップ後の拡散モデルによって予測された更新された方向と、それが指す対応する高尤度領域を示す図。

各ステップでの方向の変化は、サンプリング中に私たちが入力空間をたどるパスが曲線的であることを意味します。実際、有限の近似を行っているため、これは完全に正確ではありません。実際には区分的線形なパスです。しかし、ステップ数を無限大にすると、曲線になります。この曲線上の各点における予測方向は接線方向に対応します。以下の図は、この曲線がどのように見えるかの様式化されたバージョンです。

［画像: 無限のサンプリングステップ（破線赤線）で入力空間をたどる可能性のある曲線の様式化されたバージョンを示す図。］

図：無限のサンプリングステップ（破線赤線）で入力空間をたどる可能性のある曲線の様式化されたバージョンを示す図。

目的を持って入力空間を移動する

［画像］

多数の拡散サンプリングアルゴリズムが開発され、入力空間をより迅速に移動し、特定のレベルの出力品質を達成するために必要なサンプリングステップ数を削減しています。それらすべてをここにリストアップしようとするのは絶望的な試みでしょうが、それらの背後にある多くのアイデアが勾配ベース最適化で使用されるテクニックを模倣していることを示すために、いくつかのアルゴリズムを強調したいと思います。

拡散サンプリングに関する非常に一般的な質問は、DDPM1や確率微分方程式（SDE）ソルバー2に基づくサンプリングアルゴリズムのように、各ステップでノイズを注入すべきかどうかということです。Karrasらは彼らの「インスタントクラシック」論文（セクション3と4を参照）でこの問題を広範に研究し、確率性の導入の主な効果はエラー修正であると見出しました。拡散モデルの予測は近似的であり、ノイズは多くのサンプリングステップにわたる近似エラーの蓄積を防ぐのに役立ちます。最適化の文脈では、確率的勾配降下法（SGD）におけるノイズの正則化効果はよく研究されているため、これはおそらく驚くことではありません。

しかし、一部のアプリケーションでは、ノイズ分布からのサンプルとデータ分布からのサンプルの間の決定論的なマッピングが必要なため、各サンプリングステップでランダム性を注入することは許容されません。DDIM4やODEベースのアプローチ2のようなサンプリングアルゴリズムはこれを可能にします（私は以前にこの魔法の偉業について書きました、またこれが拡散モデルとフローベースモデルをどのように結びつけるかについてもです）。これが役立つ例としては、蒸留の文脈での教師モデル（次のセクションを参照）があります。その場合、サンプリングステップ数を増やさずに近似エラーを削減するために、他のテクニックを使用できます。

そのようなテクニックの1つは、高階メソッドの使用です。微分方程式を解くためのHeunの2次法は、2回のモデル評価を必要とするODEベースのサンプラーをもたらし、それを使用して更新方向の推定値を改善します5。これにより、各サンプリングステップのコストは約2倍になりますが、総関数評価回数という点では依然として有利なトレードオフになる可能性があります3。

このアイデアの別のバリアントは、モデルに高階スコア関数を予測させることです。これは、例えばモデルが方向と曲率の両方を推定することを意味します。これらの推定値は、曲率が低い領域ではより速く移動し、それ以外の場所では適切に減速するために使用できます。GENIE6はそのような手法の1つであり、高価な2次勾配計算を小さなニューラルネットワークに蒸留して、追加コストを実用的なレベルに削減します。

最後に、サンプリングステップ間で情報を集約することによって、高階情報を模倣できます。これは、勾配ベース最適化におけるモーメンタムの使用に非常に似ており、これも曲率に応じて加速と減速を可能にしますが、明示的に2階の量を推定する必要はありません。微分方程式解決の文脈では、このアプローチは通常マルチステップ法と呼ばれ、このアイデアは多くの拡散サンプリングアルゴリズム7 8 9 10にインスピレーションを与えています。

サンプリングアルゴリズムの選択に加えて、更新を計算する時間ステップの間隔をどのように設定するかを選択することもできます。デフォルトでは、これらのステップは全範囲にわたって均一に配置されます（np.linspaceを考えてください）。しかし、ノイズスケジュールはしばしば非線形であるため（つまり、(\sigma_t) が (t) の非線形関数であるため）、対応するノイズレベルは結果として非線形に配置されます。しかし、サンプリングステップの間隔をノイズスケジュールとは別に調整するハイパーパラメータとして扱う（または、サンプリング時にノイズスケジュールを変更する）ことで、利益が得られる可能性があります。時間ステップを賢く配置することで、指定されたステップ予算での結果の品質を向上させることができます3。

［画像］

広義には、ニューラルネットワークの文脈では、蒸留とは、あるニューラルネットワークの出力を模倣するように別のニューラルネットワークを訓練すること11を指します。前者は生徒、後者は教師と呼ばれます。通常、教師は事前に訓練されており、その重みは固定されています。拡散モデルに適用すると、興味深いことが起こります。生徒と教師のネットワークがアーキテクチャに関して同一であっても、生徒は教師が訓練されていたときよりも著しく速く収束します。

これがなぜ起こるのかを理解するために、拡散モデルのトレーニングには、データセットからの例 (\mathbf{x}_0) を使用してネットワークを監視することが含まれます。これに様々な量のノイズを加えてネットワーク入力 (\mathbf{x}_t) を作成します。しかし、ネットワークが (\mathbf{x}_0) を正確に予測することを期待するのではなく、実際に私たちが望むのは、それが (\mathbb{E}\left[,\mathbf{x}_0 \mid \mathbf{x}_t \right], すなわちデータ分布に対する条件付き期待値を予測することです。これを理解するには、この記事のセクション1の最初の図を再確認する価値があります。私たちは (\mathbf{x}_0) でモデルを監視しますが、これはモデルに予測させたいものではありません。私たちが実際に望むのは、高尤度領域の重心を指す方向を予測させることです。(\mathbf{x}_0) はその代表的なサンプルにすぎません。私は以前、拡散に関する様々な視点について議論した際にこれを言及しました。これは、トレーニングが進むにつれて、重み更新が常にモデルの重みを異なる方向に引っ張ることを意味し、収束を遅らせます。

拡散モデルを蒸留する場合、ゼロからトレーニングするのではなく、教師は (\mathbb{E}\left[,\mathbf{x}_0 \mid \mathbf{x}_t \right], の近似を提供し、生徒はそれを模倣することを学びます。以前とは異なり、モデルを監視するために使用されるターゲットは、単一の代表的なサンプルではなく、すでに（近似的な）期待値です。その結果、蒸留損失の分散は、標準的な拡散トレーニング損失の分散と比較して大幅に減少します。後者は、トレーニング曲線がいたるところでジャンプする傾向があるのに対し、蒸留ははるかにスムーズな乗り心地を提供します。これは、両方のトレーニング曲線を図で並べてプロットすると特に明らかです。この分散削減はコストを伴うことに注意してください。教師自体が不完全なモデルであるため、分散をバイアスと交換しているのです。

しかし、分散削減だけでは、拡散モデルの蒸留がこれほど人気がある理由を説明できません。蒸留は、必要なサンプリングステップ数を削減する非常に効果的な方法でもあります。これは、サンプリングアルゴリズムを変更するよりも、この点ではるかに効果的であるように思われますが、追加のモデルトレーニングが必要なため、当然ながら初期コストも高くなります。

拡散蒸留には多くのバリアントがあります。以下にそのいくつかを簡潔に要約しようとします。これは網羅的な文献レビューではないことは言うまでもありません。比較的最近のサーベイ論文は、Weijian Luoによるもの（2023年4月）12ですが、この分野ではそれ以降多くの研究が登場していますので、新しいものもいくつかカバーしようと思います。もし、言及する価値があるが私がカバーしていない特定のメソッドがあると感じた場合は、コメントで知らせてください。

拡散サンプリングを単一フォワードパスに蒸留する

典型的な拡散サンプリング手順では、ニューラルネットワークをキャンバスに繰り返し適用し、その予測を使用してキャンバスを更新します。このネットワークの計算グラフを展開すると、多くのレイヤーが重みを共有する、それ自体のはるかに深いニューラルネットワークとして再解釈できます。私は以前に、この拡散の視点についてより詳細に議論しました。

蒸留は、より大きなネットワークをより小さなネットワークに圧縮するためによく使用されます。そのため、Luhman & Luhman13は、展開されたサンプリング手順に対応する、このはるかに深い教師ネットワークの出力を再現するために、はるかに小さな生徒ネットワークを訓練することに着手しました。実際、彼らが提案しているのは、単一の拡散予測ステップに使用されるのと同じアーキテクチャを持つネットワークに、サンプリング手順全体を蒸留することです。これは、最小二乗法（MSE損失）で出力を一致させることによって行われます。サンプリング手順のステップ数によっては、これはモデル圧縮の非常に極端な形態に対応する可能性があります（計算の観点から言えば、パラメータ数はもちろん同じままです）。

このアプローチは決定論的なサンプリング手順を必要とするため、彼らはDDIM4を使用します。これは、後に開発された多くの蒸留手法も従う選択です。彼らのアプローチの結果は、ノイズ分布からのサンプルを単一フォワードパスでデータ分布からのサンプルに変換する、コンパクトな生徒ネットワークです。

［画像: 拡散サンプリング手順を単一フォワードパスに蒸留する様子を示す図。］

図：拡散サンプリング手順を単一フォワードパスに蒸留する様子を示す図。

しかし、これを実践に移すと、重大な障害に遭遇します。生徒の単一のトレーニング例を得るために、教師を使用して完全な拡散サンプリング手順を実行しなければならないのですが、これは通常、トレーニング中にオンザフライで行うにはコストが高すぎます。したがって、生徒のデータセットはオフラインで事前に生成する必要があります。これは依然として高価ですが、少なくとも一度だけ実行すればよく、生成されたトレーニング例は複数のエポックで再利用できます。

学習プロセスを高速化するために、生徒を教師の重みで初期化することも役立ちます（アーキテクチャが同一であるため可能です）。これは、ほとんどの拡散蒸留手法が利用するトリックです。

この研究は拡散蒸留の説得力のある概念実証として機能しましたが、計算コストを除けば、決定論的なサンプリング手順でのエラーの蓄積と、生徒の予測の近似的な性質が、達成可能な出力品質に著しい制限を課しました。

段階的蒸留

段階的蒸留14は、必要なサンプリングステップ数を半分にする反復的なアプローチです。これは、2つの連続するサンプリングステップの出力を単一フォワードパスに蒸留することによって達成されます。前の方法と同様に、これには決定論的なサンプリング方法（論文ではDDIMを使用）と、教師モデルに使用する所定のサンプリングステップ数 (N) が必要です。

［画像: 段階的蒸留を示す図。生徒は2つのサンプリングステップの結果を1つのフォワードパスで一致させることを学習します。］

図：段階的蒸留を示す図。生徒は2つのサンプリングステップの結果を1つのフォワードパスで一致させることを学習します。

サンプリングステップ数をさらに削減するために、これを繰り返し適用できます。理論的には、手順を (\log_2 N) 回適用することで、単一ステップサンプリングまで削減できます。これは、前のアプローチのいくつかの欠点に対処します。

  • 各蒸留ステージでは、2つの連続するサンプリングステップのみが必要であり、エンドツーエンドで完全なサンプリング手順を実行するよりも大幅にコストが低くなります。したがって、トレーニング中にオンザフライで実行でき、トレーニングデータセットの事前生成は不要になります。
  • 教師モデルに使用された元のトレーニングデータセットは、利用可能であれば（または他のデータセットでも！）再利用できます。これは、学習を関連性があり興味深い入力空間の部分に集中させるのに役立ちます。
  • 1ステップまで削減することも可能ですが、手順の反復的な性質により、品質と計算コストのトレードオフが可能になります。4ステップまたは8ステップまで削減すると、蒸留による避けられない品質低下を抑えつつ、サンプリングを大幅に高速化するのに役立ちます。これは、教師モデルのサンプリングステップ数を単純に削減するよりも、はるかに優れたトレードオフを提供します（論文の図4を参照）。

補足：v予測

画像ドメインで拡散モデルをトレーニングするための最も一般的なパラメータ化では、ニューラルネットワークが標準化されたガウスノイズ変数 (\varepsilon) を予測しますが、これは段階的蒸留に問題を引き起こします。(\varepsilon) に関するMSE損失におけるノイズレベルの暗黙的な相対重み付けは、人間の視覚システムの低周波および高周波に対する感度の違いによくマッピングするため、視覚データに特に適しています。これが、これほど一般的に使用されている理由です。

ノイズのある入力 (\mathbf{x}_t) から (\varepsilon) を予測するモデルから、入力空間での予測 (\hat{\mathbf{x}}_0) を得るには、次の式を使用できます。

\[\hat{\mathbf{x}}_0 = \alpha_t^{-1} \left( \mathbf{x}_t - \sigma_t \varepsilon (\mathbf{x}_t) \right) .\

ここで、(\sigma_t) は時間ステップ (t) におけるノイズの標準偏差を表します。（分散保存拡散の場合、スケールファクターは (\alpha_t = \sqrt{1 - \sigma_t^2})、分散爆発拡散の場合、(\alpha_t = 1) です。）

ノイズレベルが高い場合、(\mathbf{x}_t) はノイズによって支配されるため、(\mathbf{x}_t) とスケールされたノイズ予測との差は非常に小さい可能性があります。しかし、この差が入力空間での予測 (\hat{\mathbf{x}}_0) を完全に決定します！これは、予測エラーが増幅される可能性があることを意味します。標準的な拡散モデルでは、これは実際には問題ではありません。なぜなら、エラーは多くのサンプリングステップを通じて修正できるからです。段階的蒸留では、これは後期の反復で問題になります。そこでは、主にノイズレベルが高い場合にモデルを評価します（単一ステップモデルの極限では、モデルは常に最も高いノイズレベルで評価されます）。

この問題は、モデルを (\mathbf{x}_0) を予測するようにパラメータ化するだけで対処できることが判明しましたが、段階的蒸留論文では新しい予測ターゲット (\mathbf{v} = \alpha_t \varepsilon - \sigma_t \mathbf{x}_0)（「速度」、セクション4と付録Dを参照）も導入しています。これは非常に良い特性を持っており、最近では蒸留アプリケーションを超えても非常に人気があります。

ガイダンス蒸留

高品質なサンプルをより少ないステップで達成することを目的とする、より高度な拡散蒸留手法に進む前に、ガイダンス蒸留を見てみる価値があります。この手法の目的は、分類子フリーガイダンス 15を使用する際に、各ステップの計算コストを削減することです。私はすでに拡散ガイダンスに特化した2つのブログ記事を費やしているので、ここでは概念を繰り返しません。もし馴染みがない場合は、まずそれらを参照してください。

分類子フリーガイダンスの使用は、サンプリングステップごとに2回のモデル評価（1つは条件付き、1つは非条件付き）を必要とします。これにより、サンプリングコストは約2倍になります。なぜなら、主なコストはモデル評価にあるからです。そのコストを回避するために、ガイダンスから生じる予測を、選択されたガイダンススケールに基づいて単一フォワードパスで直接予測するモデルに蒸留できます16。

ガイダンス蒸留はサンプリングステップ数を削減しませんが、ステップあたりの必要な計算量を約半分に削減するため、サンプリング速度は約2倍になります。他の形態の蒸留と組み合わせることも可能です。これは有用です。なぜなら、サンプリングステップ数を削減すると、実際にはガイダンスの影響が減少するからです。ガイダンスは、更新方向の繰り返し小さな調整に依存して機能します。他の蒸留手法の前にガイダンス蒸留を適用すると、ステップ数が削減される際に元の効果が維持されることを保証するのに役立ちます。

［画像: ガイダンス蒸留を示す図。分類子フリーガイダンスを使用した単一ステップのサンプリング（拡散モデルの2回のフォワードパスが必要）が、単一フォワードパスに蒸留されます。］

図：ガイダンス蒸留を示す図。分類子フリーガイダンスを使用した単一ステップのサンプリング（拡散モデルの2回のフォワードパスが必要）が、単一フォワードパスに蒸留されます。

整流フロー

拡散サンプリングが多くの小さなステップを必要とする理由を理解する1つの方法は、曲率のレンズを通して見ることです。直線でしかステップを踏めないため、ステップが大きすぎると曲線から「外れて」しまい、顕著な近似エラーにつながります。

前述のように、一部のサンプリングアルゴリズムは、ステップサイズを決定するために曲率情報を使用したり、エラーの蓄積を減らすためにノイズを注入したりすることでこれを補償します。整流フロー法17は、より抜本的なアプローチを取ります。ノイズサンプルとデータサンプルの間のこれらの曲線パスを、はるかに曲率の低い別のパスのセットに置き換えるとしたらどうでしょうか？

これは蒸留に似た手順を使用して可能ですが、目標はまったく同じではありません。蒸留はサンプル間の既存のパスのより良い/より速い近似を学習しようとしますが、リフロー手順はパスをまったく新しいパスのセットに置き換えます。「最適輸送」の意味でのコストが低いパスの新しいセットをもたらす新しいモデルを取得します。具体的には、パスはより直線的になります。通常、以前とは異なるサンプルのペアを接続します。ある意味で、ノイズからデータへのマッピングは、より直線的になるように「配線」し直されます。

［画像: リフロー手順を適用した後、データポイントx0に関連付けられた古いパスと新しいパスを示す図。新しいパスは（完全に直線的ではないものの）著しく曲率が低く、x0を以前とは異なるノイズ分布からのサンプルに接続します。］

図：リフロー手順を適用した後、データポイントx0に関連付けられた古いパスと新しいパスを示す図。新しいパスは（完全に直線的ではないものの）著しく曲率が低く、x0を以前とは異なるノイズ分布からのサンプルに接続します。

曲率が低いということは、この新しいモデルからお気に入りのサンプリングアルゴリズムを使用してサンプリングする際に、より少なく、より大きなステップを踏むことができることを意味します。それでも近似エラーを抑えつつ。しかし、それ以外にも、これは蒸留の効果を大幅に高めます。おそらく、タスクを容易にするためです。

この手順は再帰的に適用でき、さらに直線的なパスのセットを生成できます。無限回の適用後、パスは完全に直線的になるはずです。実際には、各手順が前のものを近似する新しいモデルを生成するため、エラーが急速に蓄積する可能性があるため、これはある点までしか機能しません。幸いなことに、ほとんど直線的なパスを取得するには、1回または2回の適用で十分です。

この方法は、Stable Diffusionモデル18に適用され、その後、知覚損失19を使用した蒸留ステップが続きました。結果のモデルは、単一フォワードパスで妥当なサンプルを生成します。この方法の欠点の1つは、各リフローステップでサンプルペア（データと対応するノイズ）のデータセットを生成する必要があることです。これは通常、実用的であるためにはオフラインで行う必要があります。

整合性蒸留 & TRACT

前述のように、拡散サンプリングは入力空間を曲線的にたどり、この曲線の各点では、拡散モデルが接線方向を予測します。もし、パスの終点をデータ分布側から予測できるモデルがあればどうでしょうか？これにより、パス上のどこからでも1ステップでそこにジャンプできます。その場合、曲率の程度は問題になりません。

これが整合性モデル20が行うことです。これらは拡散モデルと非常によく似ていますが、異なる種類の量を予測します。現在の位置での接線方向ではなく、パスの終点です。ある意味で、拡散モデルと整合性モデルは、ノイズとデータの間のマッピングを記述する2つの異なる方法にすぎません。おそらく、整合性モデルを拡散モデルの「積分形式」（または等価的に、拡散モデルを整合性モデルの「微分形式」）と考えるのは有用かもしれません。

［画像: 拡散モデル（グレー）と整合性モデル（青）の予測の違いを示す図。前者はパスの接線方向を予測し、後者はデータ側のパスの終点を予測します。］

図：拡散モデル（グレー）と整合性モデル（青）の予測の違いを示す図。前者はパスの接線方向を予測し、後者はデータ側のパスの終点を予測します。

整合性モデルをゼロからトレーニングすることも可能ですが（私の意見ではそれほど簡単ではありませんが、後述します）、より実用的な方法は、まず拡散モデルをトレーニングし、その後蒸留することです。このプロセスは整合性蒸留と呼ばれます。

結果のモデルは、拡散サンプリング手順を単一フォワードパスに蒸留した場合に得られるものと非常によく似ていることに注意してください。しかし、それはパスの一方の端点（ノイズ側）からもう一方の端点（データ側）へのジャンプしかできません。整合性モデルは、パス上のどこからでもデータ側の終点にジャンプできます。

パス上の任意の点をその終点にマッピングすることを学習するには、ペアになったデータが必要なため、教師モデルを使用して完全なサンプリングプロセスを実行する必要があるように思われますが、これは高価です。しかし、これはブートストラップメカニズムを使用することで回避できます。教師からの学習に加えて、生徒は自分自身からも学習します。

これは次の原則に基づいています。パス上のすべての点における整合性モデルの予測は同じであるべきです。したがって、教師を使用してパスに沿ってステップを取ると、生徒の予測は変更されないはずです。(f(\mathbf{x}_t, t)) を生徒（整合性モデル）とすると、次のようになります。

\[f(\mathbf{x}_{t - \Delta t}, t - \Delta t) \equiv f(\mathbf{x}_t, t),\

ここで、(\Delta t) はステップサイズであり、(\mathbf{x}_{t - \Delta t}) は (\mathbf{x}_t) から始まるサンプリングステップの結果であり、更新方向は教師によって与えられます。予測はパス上のすべての点で一貫性を保ちます。これが名前の由来です。これは拡散モデルにはまったく当てはまらないことに注意してください。

整合性モデルの論文と並行して、推移閉包時間蒸留（TRACT）21が段階的蒸留の改善として提案されました。これは非常に似たブートストラップメカニズムを使用しています。実装の詳細は異なります。整合性モデルが行うように、パス上の任意の点からパスの終点を予測するのではなく、TRACTは時間ステップの範囲を区間に分割し、蒸留モデルはこれらの区間の境界にあるパス上の点を予測します。

［画像: TRACTが時間ステップ範囲を区間に分割する方法を示す図。パス上の任意の点から、生徒は現在いる区間の左境界に対応する点を予測するように訓練されます。これは整合性モデルと同じターゲットですが、パス全体ではなく、非重複セグメントに個別に適用されます。］

図：TRACTが時間ステップ範囲を区間に分割する方法を示す図。パス上の任意の点から、生徒は現在いる区間の左境界に対応する点を予測するように訓練されます。これは整合性モデルと同じターゲットですが、パス全体ではなく、非重複セグメントに個別に適用されます。

段階的蒸留と同様に、これは区間をますます少なくして繰り返し適用できる手順であり、最終的には整合性モデルとほぼ同じものになります（時間ステップ範囲全体を包含する単一区間を使用する場合）。TRACTは段階的蒸留の代替として提案され、より少ない蒸留ステージで済むため、エラー蓄積の可能性が減少します。

拡散モデルは重み平均22 23から大幅に恩恵を受けることがよく知られているため、TRACTと整合性モデルの元の定式化の両方で、生徒の重みの指数移動平均（EMA）を使用して自己教師モデルが構築されます。これは、拡散モデルと並んで、蒸留プロセスで追加の教師として機能します。とはいえ、整合性モデルのより最近のイテレーション24はEMAを使用しません。

整合性モデルを改善するための別の戦略は、蒸留に代替損失関数を使用することです。例えば、通常の使用済み平均二乗誤差（MSE）の代わりに、知覚損失（LPIPS19など）を使用することです。これは、整流フロー17でも使用されてきました。

Stable Diffusionモデルを潜在整合性モデル25に蒸留する最近の研究は、説得力のある結果をもたらし、1〜4サンプリングステップで高解像度画像を生成しています。

整合性軌跡モデル26は、拡散モデルと整合性モデルの両方の一般化であり、任意の点からその前の任意の点までのパスの予測、および接線方向を可能にします。これを達成するために、それらは開始位置と終了位置を示す2つの時間ステップに条件付けられます。両方の時間ステップが同じ場合、モデルは拡散モデルのように接線方向を予測します。

BOOT：データフリー蒸留

整合性モデルが学習するように、パス上の任意の点からデータ側のパスの終点を予測する代わりに、ノイズ側のパスの終点からパス上の任意の点を予測しようとすることができます。これがBOOT27が行うことであり、ノイズとデータの間のマッピングを記述するもう1つの方法を提供します。この定式化を整合性モデルと比較すると、一方が他方の「転置」のように見えます（下の図を参照）。word2vec28を覚えている方には、skip-gramとcontinuous bag-of-words（CBoW）メソッドの関係を強く思い出させます！

［画像: 整合性蒸留（上）とBOOT（下）の生徒の入力と予測ターゲットを示す図。Gu et al. 2023の図2に基づく。］

図：整合性蒸留（上）とBOOT（下）の生徒の入力と予測ターゲットを示す図。Gu et al. 2023の図2に基づく。

整合性モデルと同様に、この定式化は、教師を使用して完全なサンプリング手順を実行する必要を回避するためのブートストラップの一形態を可能にします（したがって、名前はおそらくそれによるものです）。生徒を使用して (\mathbf{x}_t = f(\varepsilon, t)) を予測し、教師サンプリングステップを実行して (\mathbf{x}_{t - \Delta t}) を取得し、その後 (f(\varepsilon, t - \Delta t) \equiv \mathbf{x}\{t - \Delta t}) となるように生徒を訓練します。

生徒はノイズ (\varepsilon) のみを入力として取るため、蒸留を実行するためにトレーニングデータは必要ありません。これは、拡散サンプリング手順を単一フォワードパスに直接蒸留する場合も同様です。ただし、その場合、教師を使用して完全なサンプリング手順を実行することを避けることはできません。

しかし、1つの大きな注意点があります。(\mathbf{x}_t) を予測することは、実際には学習が非常に難しいことが判明しました。しかし、これには巧妙な回避策があります。(\mathbf{x}_t) を直接予測する代わりに、まず (\mathbf{x}_t = \alpha_t \mathbf{x}_0 + \sigma_t \varepsilon) の恒等式を使用して別のターゲットに変換します。(\varepsilon) が与えられているため、これを (\mathbf{x}_0 = \frac{\mathbf{x}_t - \sigma_t \varepsilon}{\alpha_t}) と書き換えることができます。これはクリーンな入力の推定値に対応します。(\mathbf{x}_t) がノイズの多い画像のように見えるのに対し、この単一ステップの (\mathbf{x}_0) 推定値はぼやけた画像のように見え、高周波コンテンツが欠けています。これはニューラルネットワークが予測するのにずっと簡単です。

(\mathbf{x}_t) を信号とノイズの混合物と見なすと、基本的に「信号」成分を抽出し、それを代わりに予測しています。同じ式を使用して、このような予測を (\mathbf{x}_t) の予測に簡単に変換できます。(\mathbf{x}_t) が入力空間を ODE によって記述できるパスをたどるのと同様に、この時間依存の (\mathbf{x}_0) 推定値も同様です。BOOTの著者は、このパスを記述する ODE を信号-ODEと呼んでいます。

元の整合性モデルの定式化（およびTRACT）とは異なり、ブートストラップ手順には指数移動平均は使用されません。エラー蓄積を減らすために、著者は教師のサンプリングステップを実行するために高階ソルバーの使用を提案しています。この方法をうまく機能させるためのもう1つの要件は、補助的な「境界損失」であり、蒸留モデルが (t = T)（つまり、最も高いノイズレベル）でうまく振る舞うことを保証します。

ニューラルオペレータによるサンプリング

ニューラルオペレータによる拡散サンプリング（DSNO; DFNOとしても知られる、頭字語はいつか変更されたようです！)29は、ノイズサンプルからデータへのパス全体を単一フォワードパスで予測できるモデルを訓練することによって機能します。入力（(\varepsilon)）とターゲット（様々な (t) における (\mathbf{x}_t)）はBOOTで蒸留された生徒モデルと同じですが、後者は一度にパス上の単一の点しか生成できません。

これは野心的です。ニューラルネットワークがノイズからデータまで、パス全体を一度に予測できるのでしょうか？ いわゆるフーリエニューラルオペレータ（FNO）30がこれを達成するために使用されます。特定のアーキテクチャ制約を課し、時間畳み込みレイヤーを追加し、フーリエ変換を使用して周波数空間で時間の関数を表現することにより、一度に任意の数の時間ステップの予測を生成できるモデルが得られます。

自然な疑問は、なぜ実際にパス全体を予測したいのか？ということです。サンプリング時には、最終的な結果、つまりデータ側のパスの終点（(t = 0)）だけが重要です。BOOTの場合、パスの他の点を予測する目的は、トレーニングに使用されるブートストラップメカニズムを可能にすることです。しかし、DSNOにはブートストラップは含まれていないため、これを実行する意味は何でしょうか？

答えは、おそらく時間畳み込みレイヤーの帰納的バイアスと、拡散モデルが学習する入力空間を通るパスの相対的な滑らかさにあります。このアーキテクチャ上の事前知識のおかげで、パスの他の点でのトレーニングは、データ側の終点での予測の品質を向上させるのにも役立ちます。つまり、単一ステップでのサンプリング時に実際に気にするパス上の唯一の点です。私はこれが唯一の理由であると100％確信しているわけではありません。もしこれがうまくいく別の説得力のある理由があれば、教えてください！

スコア蒸留サンプリング

スコア蒸留サンプリング（SDS）31は、これまで議論してきた手法とは少し異なります。高忠実度出力をより少ないステップで実現する生徒モデルを生成することでサンプリングを高速化するのではなく、この手法は画像のパラメータ化された表現の最適化を目的としています。これは、拡散モデルが、ピクセルグリッド以外の画像表現で動作できることを可能にします。たとえそれらがそのように訓練されていたとしても、それらの表現がパラメータに関して微分可能なピクセル空間出力を生成する限りは32。

これの具体的な例として、SDSは実際にはテキストから3Dへの変換を可能にするために導入されました。これは、事前訓練された画像拡散モデルをテキストプロンプトで制御される3Dモデルの生成に使用し、ランダムな2D投影に適用することによって、Neural Radiance Field（NeRF）33表現を最適化することによって達成されます（DreamFusion）。

単純に考えれば、様々な時間ステップでの拡散損失を、パラメータ化された表現によって生成されたピクセル空間出力にバックプロパゲートするだけでうまくいくはずだと考えるかもしれません。これは、ピクセル空間出力が拡散損失を最小化するように、表現パラメータに関する勾配更新をもたらし、ピクセル空間出力がより尤もらしい画像のように見えるはずです。残念ながら、この方法は、ピクセル表現に直接適用した場合でも、あまりうまくいきません。

これは主に勾配の特定の一因によるものであることが判明しました。それは拡散モデル自体のヤコビアンに対応します。このヤコビアンは、ノイズレベルが低い場合には条件付けが悪いです。この因子を完全に省略する（つまり、単位行列に置き換える）と、事態ははるかにうまく機能します。追加の利点として、拡散モデルをバックプロパゲートする必要がなくなります。必要なのはフォワードパスのみです。これは、通常の拡散サンプリングアルゴリズムと同様です！

かなりアドホックな方法で勾配を修正した後、この修正された勾配がどの損失関数に対応するかを尋ねる価値があります。これは、確率密度蒸留 34で使用されているものと同じ損失関数です。これは元々、オーディオ波形生成用の自己回帰モデルをフィードフォワードモデルに蒸留するために開発されました。ここではこの関連性について詳しく説明しません。ただし、SDSが示すモード探索行動の説明を提供していることに言及するだけです。この行動はしばしば病理を引き起こし、それを緩和するために追加の正則化損失項が必要になります。また、教師に高いガイダンススケールを使用すること（画像をサンプリングする際に通常使用する値よりも高い値）が、結果を改善するのに役立つこともわかりました。

ノイズフリースコア蒸留（NFSD）35は、より低いガイダンススケールを可能にするために勾配をさらに変更するバリアントであり、これによりサンプル品質と多様性が向上します。変分スコア蒸留サンプリング（VSD）36は、点推定ではなく、パラメータ化された表現の分布を最適化することによってSDSを改善します。これにより、高いガイダンススケールを必要としなくなります。

VSDは、サンプリングステップ数を削減することを目的とした、より伝統的な拡散蒸留戦略のコンポーネントとして使用されてきました。単一ステップの画像ジェネレータは、パラメータ化された表現の分布として簡単に再解釈できます。これにより、VSDは、元々はテキストから3Dへの変換を改善するために考案されたものであっても、この設定に容易に適用できます。画像生成の高速化のためではありません。

Diff-Instruct37は、そのようなアプリケーションと見なすことができますが、実際にはVSDと同時に公開されました。拡散モデルの知識を単一ステップのフィードフォワードジェネレータに蒸留するために、彼らは積分KLダイバージェンス（IKL）を最小化することを提案しています。これは、時間に関するKLダイバージェンスの重み付き積分です。その勾配は、教師の予測と、ジェネレータ出力で同時に訓練されている補助的な拡散モデルの予測を対比することによって推定されます。この同時トレーニングは、それにいくらかGAN38の風味を与えますが、ジェネレータと補助モデルはここでは敵対者ではないことに注意してください。SDSと同様に、IKLのジェネレータパラメータに対する勾配は、拡散モデル教師の評価のみを必要とし、それをバックプロパゲートする必要はありません。ただし、補助拡散モデルをジェネレータ出力でトレーニングするには、当然バックプロパゲーションが必要です。

分布マッチング蒸留（DMD）39は、異なる角度から非常に類似した定式化に到達します。Diff-Instructと同様に、ジェネレータ出力の同時訓練された拡散モデルが使用され、その予測は教師の予測と対比されてフィードフォワードジェネレータの勾配を取得します。これは、教師から取得したペアデータに対する知覚回帰損失（LPIPS19）と組み合わされます。これはオフラインで事前に生成されます。後者はトレーニング例の小さなサブセットにのみ適用されるため、この事前生成ステップの計算コストはそれほど高価ではありません。

敵対的蒸留

拡散モデルが画像生成分野を完全に支配する前は、敵対的生成ネットワーク（GAN）38が最高の視覚忠実度を提供していましたが、モードドロップというコストがかかりました。モデル出力の多様性は通常、トレーニングデータの多様性を反映していませんでしたが、少なくとも見た目は良好でした。つまり、多様性を品質と交換しました。それに加えて、GANは単一フォワードパスで画像を生成するため、非常に高速です。拡散モデルのサンプリングよりもはるかに高速です。

したがって、一部の研究が敵対的モデルと拡散モデルの利点を組み合わせようとしたことは驚くことではありません。これを行う方法はたくさんあります。デノイジング拡散GAN40や敵対的スコアマッチング41はその例です。

より最近の例はUFOGen42です。これは、拡散モデルの敵対的ファインチューニングを提案しており、蒸留に非常に似ていますが、厳密な意味では実際には蒸留ではありません。UFOGenは、標準的な拡散損失と敵対的損失を組み合わせます。標準的な拡散損失だけでは、条件付き期待値 (\mathbb{E}\left[,\mathbf{x}_0 \mid \mathbf{x}_t \right], を予測しようとするモデルにつながりますが、追加の敵対的損失項により、モデルはこれと異なる予測を生成し、高ノイズレベルでぼやけの少ない予測を生成できます。結果として多様性が低下しますが、サンプリングも高速化されます。ジェネレータとディスクリミネータの両方は、事前訓練された拡散モデルのパラメータから初期化されますが、この事前訓練されたモデルは、蒸留アプローチの場合のように、トレーニングターゲットを生成するために評価されるわけではありません。それにもかかわらず、これは、議論してきたほとんどの蒸留アプローチと同じ目標を達成することを意図しているため、ここに含める価値があります。

一方、敵対的拡散蒸留 43は、「真の」蒸留アプローチであり、スコア蒸留サンプリング（SDS）と敵対的損失を組み合わせています。これは、以前は純粋な敵対的テキストから画像モデルであるStyleGAN-T45にも使用されていた画像表現学習モデル、DINO44の特徴の上に構築されたディスクリミネータを使用します。結果の生徒モデルは単一ステップサンプリングを可能にしますが、結果の品質を向上させるために複数ステップでサンプリングすることもできます。この方法は、リアルタイム生成を可能にするテキストから画像システムであるSDXL Turboに使用されました。生成された画像は、入力中に更新されます。

しかし、「フリーランチ」はないのか？

［画像］

なぜ、拡散モデルが同じことを達成するために数十または数百ステップを要するのに、これらの蒸留モデルはわずか数ステップで説得力のあるサンプルを生成できるのでしょうか？ 「フリーランチはない」というのはどういうことでしょうか？

一見すると、拡散蒸留は確かに機械学習における普遍的な真実として広く考えられているものに対する反例のように見えますが、それ以上のものがあります。ある程度までは、拡散モデルのサンプリングは、モデル品質に目立ったコストなしに蒸留によってより効率的に行うことができるでしょう。しかし、ほとんどの蒸留手法がターゲットとするレジーム（つまり、1〜4サンプリングステップ）はその点をはるかに超えており、品質と速度を交換しています。蒸留は実際にはほとんどの場合「ロスのある」ものであり、生徒は教師の予測を完全に模倣することは期待できません。これにより、エラーが発生し、サンプリングステップ全体にわたって蓄積する可能性があります。あるいは、一部の手法では、蒸留の異なるフェーズ全体にわたって蓄積する可能性があります。

このトレードオフはどのようなものなのでしょうか？それは蒸留手法によって異なります。 ほとんどの手法では、モデル品質の低下は出力の知覚品質に直接影響します。蒸留モデルからのサンプルは、しばしばぼやけて見えることがあります。あるいは、細かいディテールはシャープに見えても、人間のような顔の画像で特に顕著な、それほど現実的ではないように見えることがあります。ディスクリミネータに基づく敵対的損失、またはLPIPS19のような知覚損失関数の使用は、モデル容量を、知覚的に関連性の高い信号コンテンツにさらに焦点を当てることによって、この劣化のいくつかを軽減することを目的としています。

一部の手法は、出力品質と高周波コンテンツの忠実度を驚くほど高いレベルで維持しますが、これは通常、サンプルの多様性のコストで来ます。前述の敵対的手法は、スコア蒸留サンプリングに基づく手法と同様に、暗黙的にモード探索損失関数を最適化するため、その良い例です。

したがって、蒸留がモデル品質の低下を意味する場合、拡散モデルをトレーニングしてから蒸留することは価値があるのでしょうか？ GANのような別の種類のモデルを代わりにトレーニングする方が良いのではないでしょうか？GANは蒸留を必要とせずに、最初から単一ステップジェネレータを生成します。鍵は、蒸留がこのトレードオフを制御するある程度の自由度を提供してくれることです。私たちは柔軟性を得ます。私たちは、どれだけのステップを許容できるかを選択でき、適切な手法を選択することで、どのように手抜きをするかを正確に決定できます。忠実度と多様性のどちらをより重視しますか？それはあなたの選択です！

本当に教師が必要か？

［画像］

拡散蒸留が、出力品質、多様性、サンプリング速度の点で適切なトレードオフを持つ、私たちが求めているモデルをもたらしてくれると確立したら、そもそもこのモデルに到達するために蒸留が必要だったかどうかを問う価値があります。ある意味では、蒸留を通じて特定のモデルを取得したら、それは存在証明であり、そのようなモデルが実際には実現可能であることを示していますが、それが最も効率的な方法でそのモデルに到達したことを証明するものではありません。おそらく、より短いルートがあるのではないでしょうか？ゼロからそのようなモデルをトレーニングすることは可能でしょうか？そして、教師モデルのトレーニングを完全にスキップすることはできるでしょうか？

答えは蒸留手法によります。 拡散蒸留を通じて得られる特定の種類のモデルについては、実際には蒸留をまったく必要としない代替トレーニングレシピが存在します。しかし、これらは通常、蒸留ルートほどうまく機能しません。おそらく、これはそれほど驚くべきことではありません。大きなニューラルネットワークを小さなネットワークに蒸留する場合、ゼロからその小さなネットワークをトレーニングする場合よりも優れた結果を得られることが多いことが、以前から知られています11。同じ現象がここで起こっています。なぜなら、私たちは多くのステップを持つサンプリング手順を、はるかに少ないステップを持つ手順に蒸留しているからです。これらのサンプリング手順の計算グラフを見ると、前者は後者よりもはるかに「深い」ため、私たちがやっていることは、大きなモデルを小さなモデルに蒸留することに非常によく似ています。

蒸留かゼロからのトレーニングかを選択できる例の1つは、整合性モデルです。それらを導入した論文20は、整合性蒸留と整合性トレーニングの両方を説明しています。後者はうまく機能するためにいくつかのトリックを必要とします。ハイパーパラメータのスケジュールを作成して一種の「カリキュラム」を作成するなど、より複雑であると言えます。

データとノイズの間の迷宮をたどる

［画像］

拡散モデルのトレーニングに関する興味深い視点の1つは、それが分布間の最適輸送マップを明らかにする方法を提供することです46。確率フローODE定式化2を通じて、拡散モデルがノイズとデータの間の双射を学習することがわかります。そして、このマッピングがある意味で近似的に最適であることが判明します。

これはまた、同様のデータで訓練された異なる拡散モデルが類似のマッピングを学習するという観察結果を説明しています。それらはすべて同じ最適値を近似しようとしています！私はしばらく前にこれについてツイートしました（Xしました？）:

最近の拡散モデルを単一パスモデルに蒸留する研究の多さから、拡散モデルのトレーニングを一種の最適輸送問題の解決として考えてきました🚐 (1/6)

— Sander Dieleman (@sedielem) 2023年12月5日

これまでのところ、この最適マッピングを近似する最もシンプルで効果的な（つまりスケーラブルな）方法であるように思われますが、唯一の方法ではありません。整合性トレーニングは説得力のある代替戦略を表します。これにより、まだ発見されていない他のアプローチが何であるか、そして拡散モデルのトレーニングよりもシンプルで、統計的に効率的な方法を見つけられるかどうか疑問に思います。

これらの手法の一部間のもう1つの興味深い関連性は、曲率をより詳しく調べることで見つけることができます。ノイズとデータ分布からのサンプルを接続するパスは、拡散モデルのトレーニングによって明らかになるものは曲線的である傾向があります。これが、サンプリング時にそれらを正確に近似するために多くの離散ステップが必要な理由です。

この問題を回避するためのいくつかの方法を議論しました。整合性モデル20 21は、モデルの予測ターゲットを、現在の位置での接線方向からデータ側の終点に変更することによって回避します。一方、整流フロー17は、曲線パスを、はるかに直線的なパスのセットに置き換えます。しかし、完全に直線的なパスの場合、接線方向は実際には終点を指します！つまり、完全に直線的なパスの極限ケースでは、整合性モデルと拡散モデルは同じものを予測し、互いに区別がつかなくなります。

その観察は実用的ですか？おそらくそうではありません。単に巧妙なつながりです。しかし、分布間の決定論的なマッピングを理解し、それを大規模に明らかにする方法、そしてそれらをパラメータ化し表現するさまざまな方法をさらに深く理解することは価値があると思います。これは、拡散蒸留だけでなく、より広範な意味での反復的洗練による生成モデリングにおける革新の肥沃な土壌だと思います。

締めくくりの考察

［画像］

冒頭で述べたように、これは拡散蒸留に関するかなりハイレベルな扱いであり、なぜそれほど多くの異なる方法があるのかを説明するつもりでした。結局、これらの手法間のすべてのつながりについて話すことは、手法自体を説明せずに難しいため、少し深掘りすることになりました。このテーマについて調べ、物事を簡潔に説明しようとすることで、私は多くのことを学びました。機械学習研究の特定のテーマ（または実際には他の何でも）について学びたい場合は、それについてブログ記事を書くことを強くお勧めします。

まとめとして、いくつかのパターンとトレンドを特定したいと思います。拡散蒸留手法には多種多様なものがありますが、頻繁に戻ってくる共通のトリックやアイデアがいくつかあります。

  • 教師からターゲットを取得するために決定論的なサンプリングアルゴリズムを使用することは、ほとんどすべての手法が依存していることです。DDIM4が人気ですが、より高度な手法（例：高階手法）も選択肢です。
  • 生徒ネットワークのパラメータは通常、教師から初期化されます。これは収束を加速するだけでなく、一部の手法ではそれが機能するために不可欠です。教師と生徒のアーキテクチャが同一であることが多いため、これは可能です。これは、識別モデルの蒸留とは異なります。
  • いくつかの手法は、低レベルの知覚品質に対する蒸留の悪影響を軽減するために、LPIPS19のような知覚損失を使用しています。
  • ブートストラップ、つまり生徒が自分自身から学習することは、教師からターゲットを取得するために完全なサンプリングアルゴリズムを実行する必要を回避するための便利なトリックです。場合によっては、生徒のパラメータの指数移動平均を使用することが役立つことがわかっていますが、これはそれほど明確ではありません。

蒸留は他のモデリングの選択と相互作用する可能性があります。 重要な例の1つは分類子フリーガイダンス 15です。これは、多くのサンプリングステップが存在することに暗黙的に依存しています https://sander.ai/2023/08/28/geometry.html#closing-thoughts。ガイダンスは、拡散モデルによって予測される入力空間の方向を変更することによって機能し、サンプリングステップ数が少ない場合、この効果は必然的に減少します。一部の手法では、生徒はもはや入力空間の方向を予測しないため、蒸留後のガイダンスの適用は実際にはもはや意味がありません。幸いなことに、ガイダンス蒸留16は、この影響を軽減するために使用できます。

もう1つの例は潜在拡散 47です。潜在空間で訓練された拡散モデルに蒸留を適用する場合、対処すべき重要な問題は、損失を潜在表現またはピクセルに適用すべきかどうかです。例として、敵対的拡散蒸留（ADD）論文43は、安定性を向上させるために、ピクセル空間で蒸留損失を計算することを明示的に提案しています。

まず問題を可能な限り最善の方法で解決し、次に許容可能なトレードオフをもたらすショートカットを探すという手順は、一般的に機械学習において非常に効果的です。拡散蒸留はその典型的な例です。フリーランチはまだありませんが、拡散蒸留は意図的に手抜きをすることを可能にし、それは非常に価値があります！

この投稿を学術的な文脈で引用したい場合は、このBibTeXスニペットを使用できます。

@misc{dieleman2024distillation,
  author = {Dieleman, Sander},
  title = {The paradox of diffusion distillation},
  url = {https://sander.ai/2024/02/28/paradox.html},
  year = {2024}
}

謝辞

モデルとして再びBundle the bunnyに感謝し、この写真の使用許可をいただいたkipplyに感謝します。Google DeepMindのEmiel Hoogeboom、Valentin De Bortoli、Pierre Richemond、Andriy Mnih、そして同僚の皆さんに、拡散モデルとその先に関する私の考えを形成し続けている様々な議論に感謝します！

参考文献

  1. Ho, Jain, Abbeel, “Denoising Diffusion Probabilistic Models”, 2020. ↩ ↩2

  2. Song, Sohl-Dickstein, Kingma, Kumar, Ermon and Poole, “Score-Based Generative Modeling through Stochastic Differential Equations”, International Conference on Learning Representations, 2021. ↩ ↩2 ↩3

  3. Karras, Aittala, Aila, Laine, “Elucidating the Design Space of Diffusion-Based Generative Models”, Neural Information Processing Systems, 2022. ↩ ↩2 ↩3

  4. Song, Meng, Ermon, “Denoising Diffusion Implicit Models”, International Conference on Learning Representations, 2021. ↩ ↩2 ↩3

  5. Jolicoeur-Martineau, Li, Piché-Taillefer, Kachman, Mitliagkas, “Gotta Go Fast When Generating Data with Score-Based Models”, arXiv, 2021. ↩

  6. Dockhorn, Vahdat, Kreis, “GENIE: Higher-Order Denoising Diffusion Solvers”, Neural Information Processing Systems, 2022. ↩

  7. Liu, Ren, Lin, Zhao, “Pseudo Numerical Methods for Diffusion Models on Manifolds”, International Conference on Learning Representations, 2022. ↩

  8. Zhang, Chen, “Fast Sampling of Diffusion Models with Exponential Integrator”, International Conference on Learning Representations, 2023. ↩

  9. Lu, Zhou, Bao, Chen, Li, Zhu, “DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps”, Neural Information Processing Systems, 2022. ↩

  10. Lu, Zhou, Bao, Chen, Li, Zhu, “DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models”, arXiv, 2022. ↩

  11. Hinton, Vinyals, Dean, “Distilling the Knowledge in a Neural Network”, NeurIPS Deep Learning Workshop, 2014. ↩ ↩2

  12. Luo, “A Comprehensive Survey on Knowledge Distillation of Diffusion Models”, arXiv, 2023. ↩

  13. Luhman, Luhman, “Knowledge Distillation in Iterative Generative Models for Improved Sampling Speed”, arXiv, 2021. ↩

  14. Salimans, Ho, “Progressive Distillation for Fast Sampling of Diffusion Models”, International Conference on Learning Representations, 2022. ↩

  15. Ho, Salimans, “Classifier-Free Diffusion Guidance”, Neural Information Processing Systems, 2021. ↩ ↩2

  16. Meng, Rombach, Gao, Kingma, Ermon, Ho, Salimans, “On Distillation of Guided Diffusion Models”, Computer Vision and Pattern Recognition, 2023. ↩ ↩2

  17. Liu, Gong, Liu, “Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow”, International Conference on Learning Representations, 2023. ↩ ↩2 ↩3

  18. Liu, Zhang, Ma, Peng, Liu, “InstaFlow: One Step is Enough for High-Quality Diffusion-Based Text-to-Image Generation”, arXiv, 2023. ↩

  19. Zhang, Isola, Efros, Shechtman, Wang, “The Unreasonable Effectiveness of Deep Features as a Perceptual Metric”, Computer Vision and Pattern Recognition, 2018. ↩ ↩2 ↩3 ↩4 ↩5

  20. Song, Dhariwal, Chen, Sutskever, “Consistency Models”, International Conference on Machine Learning, 2023. ↩ ↩2 ↩3

  21. Berthelot, Autef, Lin, Yap, Zhai, Hu, Zheng, Talbott, Gu, “TRACT: Denoising Diffusion Models with Transitive Closure Time-Distillation”, arXiv, 2023. ↩ ↩2

  22. Song, Ermon, “Improved Techniques for Training Score-Based Generative Models”, Neural Information Processing Systems, 2020. ↩

  23. Karras, Aittala, Lehtinen, Hellsten, Aila, Laine, “Analyzing and Improving the Training Dynamics of Diffusion Models”, arXiv, 2023. ↩

  24. Song, Dhariwal, “Improved Techniques for Training Consistency Models”, International Conference on Learnign Representations, 2024. ↩

  25. Luo,Tan, Huang, Li, Zhao, “Latent Consistency Models: Synthesizing High-Resolution Images with Few-Step Inference”, arXiv, 2023. ↩

  26. Kim, Lai, Liao, Murata, Takida, Uesaka, He, Mitsufuji, Ermon, “Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion”, International Conference on Learning Representations, 2024. ↩

  27. Gu, Zhai, Zhang, Liu, Susskind, “BOOT: Data-free Distillation of Denoising Diffusion Models with Bootstrapping”, arXiv, 2023. ↩

  28. Mikolov, Chen, Corrado, Dean, “Efficient Estimation of Word Representations in Vector Space”, International Conference on Learning Representation, 2013. ↩

  29. Zheng, Nie, Vahdat, Azizzadenesheli, Anandkumar, “Fast Sampling of Diffusion Models via Operator Learning”, International Conference on Machine Learning, 2023. ↩

  30. Li, Kovachki, Azizzadenesheli, Liu, Bhattacharya, Stuart, Anandkumar, “Fourier neural operator for parametric partial differential equations”, International Conference on Learning Representations, 2021. ↩

  31. Poole, Jain, Barron, Mildenhall, “DreamFusion: Text-to-3D using 2D Diffusion”, arXiv, 2022. ↩

  32. Mordvintsev, Pezzotti, Schubert, Olah, “Differentiable Image Parameterizations”, Distill, 2018. ↩

  33. Mildenhall, Srinivasan, Tancik, Barron, Ramamoorthi, Ng, “NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis”, European Conference on Computer Vision, 2020. ↩

  34. Van den Oord, Li, Babuschkin, Simonyan, Vinyals, Kavukcuoglu, van den Driessche, Lockhart, Cobo, Stimberg, Casagrande, Grewe, Noury, Dieleman, Elsen, Kalchbrenner, Zen, Graves, King, Walters, Belov and Hassabis, “Parallel WaveNet: Fast High-Fidelity Speech Synthesis”, International Conference on Machine Learning, 2018. ↩

  35. Katzir, Patashnik, Cohen-Or, Lischinski, “Noise-free Score Distillation”, International Conference on Learning Representations, 2024. ↩

  36. Wang, Lu, Wang, Bao, Li, Su, Zhu, “ProlificDreamer: High-Fidelity and Diverse Text-to-3D Generation with Variational Score Distillation”, Neural Information Processing Systems, 2023. ↩

  37. Luo, Hu, Zhang, Sun, Li, Zhang, “Diff-Instruct: A Universal Approach for Transferring Knowledge From Pre-trained Diffusion Models”, Neural Information Processing Systems, 2023. ↩

  38. Goodfellow, Pouget-Abadie, Mirza, Xu, Warde-Farley, Ozair, Courville and Bengio, “Generative Adversarial Nets”, Neural Information Processing Systems, 2014. ↩ ↩2

  39. Yin, Gharbi, Zhang, Shechtman, Durand, Freeman, Park, “One-step Diffusion with Distribution Matching Distillation”, arXiv, 2023. ↩

  40. Xiao, Kreis, Vahdat, “Tackling the Generative Learning Trilemma with Denoising Diffusion GANs”, International Conference on Learning Representations, 2022. ↩

  41. Jolicoeur-Martineau, Piché-Taillefer, Tachet des Combes, Mitliagkas, “Adversarial score matching and improved sampling for image generation”, International Conference on Learning Representations, 2021. ↩

  42. Xu, Zhao, Xiao, Hou, “UFOGen: You Forward Once Large Scale Text-to-Image Generation via Diffusion GANs”, arXiv, 2023. ↩

  43. Sauer, Lorenz, Blattmann, Rombach, “Adversarial Diffusion Distillation”, arXiv, 2023. ↩ ↩2

  44. Caron, Touvron, Misra, Jégou, Mairal, Bojanowski, Joulin, “Emerging Properties in Self-Supervised Vision Transformers”, International Conference on Computer Vision, 2021. ↩

  45. Sauer, Karras, Laine, Geiger, Aila, “StyleGAN-T: Unlocking the Power of GANs for Fast Large-Scale Text-to-Image Synthesis”, International Conference on Machine Learning, 2023. ↩

  46. Khrulkov, Ryzhakov, Chertkov, Oseledets, “Understanding DDPM Latent Codes Through Optimal Transport”, International Conference on Learning Representations, 2023. ↩

  47. Rombach, Blattmann, Lorenz, Esser, Ommer, “High-Resolution Image Synthesis with Latent Diffusion Models”, Computer Vision and Pattern Recognition, 2022. ↩