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

拡散モデルは、高次元分布からのデータ生成という困難なタスクを、それぞれより簡単な多数のノイズ除去タスクに分割します。私たちは、それらを一度に1つのタスクだけを解決するように訓練します。サンプリングするには、多くの予測を逐次的に行います。この反復的な洗練こそが、それらの力の源泉です。…あるいは、そうではないのでしょうか？拡散モデルに関する最近の多くの論文は、必要なサンプリングステップ数を減らすことに焦点を当てており、一部の研究ではシングルステップでのサンプリングさえ目指しています。これは、物事を多くの簡単なステップに分割することが、そもそもこれらのモデルがうまく機能する理由であるとされるのに、逆行しているように思えます。この記事では、拡散モデルから良好な結果を得るために必要なサンプリングステップ数を減らすことができる様々な方法を詳しく見ていきます。特に、様々な形態の蒸留に焦点を当てます。これは、別のモデル（教師）の予測によって監督される新しいモデル（学生）を訓練する実践です。拡散モデルの様々な蒸留手法は、非常に説得力のある結果を生み出しています。書き始めたときは比較的ハイレベルなものにするつもりでしたが、拡散モデルの蒸留は少しニッチなトピックであるため、いくつかのことを詳細に説明せざるを得ず、ディープダイブとなりました。以下は目次です。特定のセクションに直接ジャンプするにはクリックしてください。拡散サンプリング：慎重に進む入力空間を目的を持って移動する拡散蒸留拡散サンプリングを単一のフォワードパスに蒸留する段階的蒸留ガイダンス蒸留修正フロー一貫性蒸留＆TRACT BOOT：データフリー蒸留ニューラルオペレーターを用いたサンプリングスコア蒸留サンプリング敵対的蒸留しかし、「無料のランチ」はどうでしょうか？本当に教師が必要なのでしょうか？データとノイズの間の迷路をマッピングする終わりに考える謝辞参考文献拡散サンプリング：慎重に進むまず第一に、なぜ拡散モデルから良好な結果を得るのに多くのステップが必要なのでしょうか？様々な方法が、出力の品質を損なうことなく、あるいは少なくともそれほど損なうことなく、このステップ数を削減できる理由を理解することは価値があります。拡散モデルにおけるサンプリングステップは、以下のことで構成されます。ノイズを除去するために、あるいは同等に、入力をデータ分布の下でより確率的にするために、入力空間でどの方向に移動すべきかを予測する。その方向に小さなステップを取る。サンプリングアルゴリズムによっては、少しノイズを追加したり、更新方向を計算するために、より高度なメカニズムを使用したりする場合があります。私たちは小さなステップしか取りません。なぜなら、この予測された方向は局所的にしか意味がないからです。それは、データ分布の下での確率が高い入力空間の領域を指しますが、特定のデータポイントを指すわけではありません。したがって、もし大きなステップを取った場合、その高確率領域の中心に到達することになりますが、それは必ずしもデータ分布の代表的なサンプルではありません。それを大まかな推定と考えてください。もしこれが直感的でないと感じるなら、あなたは一人ではありません！高次元空間における確率分布は、しばしば直感に反して振る舞います。これは私が以前、拡散ガイダンスの幾何学に関する詳細なブログ記事で書いたことです。具体的には、画像ドメインでは、予測された方向に大きなステップを取ると、入力に多くのノイズがある場合、ぼやけた画像が得られる傾向があります。これは、それが基本的に多くの可能な画像の平均に対応するためです。（議論のために、サンプリングアルゴリズムの一部として後で追加される可能性のあるノイズは意図的に無視しています。）別の見方をすれば、ノイズは高周波情報を不明瞭にします。これは、シャープな特徴や細かいディテールに対応します（これも以前書いたことがあります）。この高周波情報に関する不確実性は、すべての可能性が混ざり合った予測につながり、結果として高周波情報が完全に失われます。予測された方向の局所的な妥当性は、無限小のステップしか取れないことを意味します。そして、新しい方向を決定するためにモデルを再評価します。もちろん、これは実用的ではないので、有限ですが小さなステップを取ります。これは、パラメータ空間における機械学習モデルの勾配ベースの最適化の働き方に非常に似ていますが、ここでは入力空間で動作しています。モデルのトレーニングと同様に、もし取るステップが大きすぎると、最終結果の品質が低下します。以下は、2次元で入力空間を表す図です。 ext{x}_t は時間ステップ t におけるノイズの多い入力を表し、ここではデータ分布から描画されたクリーンな画像 ext{x}_0 にノイズを追加することによって構築しました。また、入力がより確率的になるために移動すべき方向（拡散モデルによって予測された）も示されています。これは、高確率領域の中心である ext{x}_0 の推定値に向かっています。高確率領域と、拡散モデルによって予測された方向を示し、その領域の中心を指している図。（非常に高次元の空間を2Dで表現することに関する注意点については、以前の拡散ガイダンスの幾何学に関するブログ記事の最初のセクションを参照してください！）この方向にステップを進め、いくらかのノイズを追加すると（例えば、DDPM1サンプリングアルゴリズムで行うように）、 ext{x}_{t-1} に到達します。これは、わずかにノイズの少ない入力画像に対応します。予測された方向は、前のサンプリングステップによって不確実性が解消されたため、より小さく、「より具体的」な高確率領域を指すようになります。これは以下の図に示されています。単一のサンプリングステップの後、拡散モデルによって予測された更新された方向と、それが指す対応する高確率領域を示している図。各ステップでの方向の変化は、サンプリング中に私たちが入力空間をたどるパスが曲線であることを意味します。実際、有限の近似を行っているため、それは完全に正確ではありません。それは実際には区分的線形パスです。しかし、ステップ数を無限大にすると、曲線になります。この曲線上の各点における予測された方向は、接線の方向に相当します。無限ステップで入力空間をたどる可能性のある曲線の様式化されたバージョンが以下の図に示されています。無限ステップ（破線の赤い曲線）で入力空間をたどる可能性のある曲線の様式化されたバージョンを示している図。入力空間を目的を持って移動する多数の拡散サンプリングアルゴリズムが、入力空間をより迅速に移動し、特定のレベルの出力品質を達成するために必要なサンプリングステップ数を削減するために開発されてきました。ここにそれらすべてをリストアップしようとするのは絶望的な試みですが、それらの背後にある多くのアイデアが勾配ベースの最適化で使用されるテクニックを模倣していることを示すために、いくつかのアルゴリズムを強調したいと思います。拡散サンプリングに関する非常に一般的な質問は、DDPM1や確率微分方程式（SDE）ソルバー2に基づいたサンプリングアルゴリズムのように、各ステップでノイズを注入すべきかどうかということです。Karrasら3は、この質問を広範囲に研究しており（彼らの「インスタントクラシック」論文のセクション3と4を参照）、確率性の導入の主な効果はエラー訂正であると結論付けています。拡散モデルの予測は近似的であり、ノイズはこれらの近似エラーが多くのサンプリングステップにわたって蓄積するのを防ぐのに役立ちます。最適化の文脈では、確率的勾配降下法（SGD）におけるノイズの正則化効果はよく研究されているため、これはおそらく驚くべきことではありません。しかし、一部のアプリケーションでは、ノイズ分布からのサンプルとデータ分布からのサンプルの間に決定論的なマッピングが必要なため、各サンプリングステップでランダム性を注入することは許容されません。DDIM4やODEベースのアプローチ2のようなサンプリングアルゴリズムは、これを可能にします（以前に書いたことがあります。