拡散モデル推論スケーリング新パラダイム

拡散モデル推論のスケーリング：新たなパラダイム

近年、大規模言語モデル(LLMs)の分野では、推論時のスケーリングが有効であることが示されています。o1、o3、DeepSeek R1、QwQ、Step Reasoner miniなどのモデルは、推論時の計算量を増やすことで性能が大幅に向上することが実証されています。この成功を受け、拡散モデルにも同じ原理が適用できるのかという疑問が生じました。

ニューヨーク大学のXie Saining氏率いるチームは、この疑問を探求しました。一般的な検索フレームワークを用いた体系的な調査の結果、拡散モデルにおいても推論時のスケーリングが有効であることが判明しました。推論時の計算量を増やすことで、生成サンプルの質が大幅に向上することが示されました。さらに、画像の複雑さから、さまざまなアプリケーションシナリオに合わせてフレームワーク内のコンポーネントを柔軟に組み合わせることが可能であることがわかりました。

主な発見

推論時のスケーリングは拡散モデルに有効: 推論時に多くの計算リソースを割り当てることで、より高品質のサンプルが得られます。
コンポーネントの組み合わせの柔軟性: このフレームワークにより、さまざまなアプリケーションに対応できる多様なコンポーネント構成が可能です。
ノイズ除去ステップの先へ: 研究では、単にノイズ除去ステップを増やすだけでなく、サンプリング中のより良いノイズを探索することが、NFE（関数評価数）をスケーリングするための別の次元であることが示唆されています。
二つの設計軸: フレームワークは、以下の二つの主要な設計軸に焦点を当てています。
- 検証器（Verifiers）: 検索プロセス中にフィードバックを提供します。
- アルゴリズム（Algorithms）: より良いノイズ候補を見つけます。

研究方法論

研究チームは、様々なユースケースをシミュレートするために、検証器に対して3つの異なるシナリオを検討しました。

最終評価に関する特権情報が利用可能なシナリオ
生成を導くための条件付き情報が利用可能なシナリオ
追加情報が利用できないシナリオ

アルゴリズムについては、以下の3つを調査しました。

ランダムサーチ: 固定された候補セットから最適なものを選択します。
ゼロオーダーサーチ: 検証器のフィードバックを用いて、ノイズ候補を反復的に改善します。
パスサーチ: 検証器のフィードバックを用いて、拡散サンプリング軌跡を反復的に改善します。

この研究では、最初に比較的シンプルなImageNetクラス条件付き生成設定でこれらの設計を検討しました。その後、大規模なテキスト条件付き生成にこれらの設計を適用し、提案されたフレームワークを評価しました。

推論時間のスケーリング

この論文では、拡散モデルにおける推論時間をスケーリングするためのフレームワークを提案し、最適なサンプリングノイズの探索として課題を捉えています。このプロセスには、2つのコアコンポーネントが含まれます。

検証器（Verifiers）: 生成されたサンプルの質を評価する事前学習済みモデルです。生成されたサンプルとオプションの条件を入力として受け取り、スカラー値をスコアとして出力します。
アルゴリズム（Algorithms）: 検証器のスコアを使用して、より良い候補サンプルを見つけるアルゴリズムです。関数は、検証器(V)、事前学習済みの拡散モデル(_θ)、生成されたサンプルと条件のセットを入力として受け取り、最適な初期ノイズを出力します。

総推論予算は、ノイズ除去ステップと探索コストの両方を含む、関数評価の総数(NFE)によって測定されます。

探索検証器

研究者らは、最初に、選択されたサンプルの最終評価に関する完全な情報を持つオラクル検証器から始めました。ImageNetの場合、これにはFIDやISなどのメトリクスが含まれていました。次に、CLIPやDINOなどの、よりアクセスしやすい事前学習済みモデルを教師あり検証器として検討しました。これらのモデルはサンプルを分類するために使用され、クラスラベルに対応する最も高いロジットを持つサンプルが選択されました。

しかし、これらの分類器は点単位で動作するため、FIDスコアの目的と部分的にしか一致しないことが観察されました。これにより、計算量が増加するにつれて、サンプル分散の減少とモード崩壊が発生しました。この現象は「検証器ハッキング」と呼ばれ、ランダムサーチアルゴリズムの制約のない探索空間によって加速されました。

興味深いことに、検証器は効果的に探索を導くために必ずしも条件付き情報を必要としないことが研究でわかりました。DINO/CLIP分類器からのロジットと、低ノイズレベルでのx予測と最終的なクリーンサンプルとの間の特徴空間のコサイン類似度との間に強い相関関係が観察されました。これにより、追加の条件付き情報を必要とせず、依然として効果的なスケーリング動作を示す自己教師あり検証器の使用につながりました。

探索アルゴリズム

検証器ハッキングを緩和するために、研究者らは候補サンプルを徐々に最適化する、より洗練された探索アルゴリズムを検討しました。これには、以下のゼロオーダーサーチ法が含まれます。

ランダムなガウスノイズをピボットポイントとして開始します。
ピボットポイントの近傍でN個の候補を見つけます。
ODEソルバーを通じて候補を実行し、サンプルと検証器スコアを取得します。
最適な候補でピボットポイントを更新し、ステップ1〜3を繰り返します。

また、サンプリング軌跡に沿って探索する可能性を検討するパスサーチアルゴリズムも調査しました。

N個の初期ノイズサンプルをサンプリングし、ODEソルバーをノイズレベルσまで実行します。
各サンプルにノイズを追加し、順方向ノイズ処理をシミュレートします。
各ノイズ付きサンプルでODEソルバーを実行し、検証器スコアに基づいて上位N個の候補を保持し、ODEソルバーがσ=0に達するまで繰り返します。
残りのN個のサンプルをランダムに探索し、最適なものを保持します。

ゼロオーダーとパスサーチアルゴリズムはどちらも、ランダムサーチと比較して強い局所性を維持します。

テキストから画像生成へのスケーリング

研究チームは、大規模なテキストから画像へのタスクにおける検索フレームワークのスケーリング能力を検証しました。DrawBenchおよびT2I-CompBenchデータセットを評価に使用し、バックボーンとしてFLUX.1-devモデルを使用しました。また、Aesthetic Score Predictor、CLIPScore、ImageRewardなどの教師あり検証器の選択肢も拡大しました。さらに、これらの3つの検証器を組み合わせることで、検証器アンサンブルを作成しました。

分析：検証器とタスクの整合性

この研究では、さまざまなデータセットにおける、さまざまな検証器とアルゴリズムの組み合わせの結果を比較しました。DrawBenchでは、すべての検証器を使用すると、一般的にサンプルの質が向上することがわかりました。しかし、AestheticとCLIP検証器を単独で使用すると、それぞれのバイアスに過剰適合し、互いに悪影響を及ぼす可能性があることが観察されました。これは、評価の焦点のミスマッチに起因します。Aesthetic Scoreは視覚的な質に焦点を当て、高度に様式化された画像を好む傾向がある一方、CLIPは視覚とテキストの整合性を優先し、視覚的な質を犠牲にすることがあります。

一部の検証器は特定のタスクに適しており、検証器の有効性はタスク要件との整合性に依存することが指摘されました。

アルゴリズムのパフォーマンス

3つの検索アルゴリズム（ランダム、ゼロオーダー、パス）はすべて、DrawBenchでのサンプリング品質を効果的に向上させました。しかし、ランダムサーチは他の2つの方法の局所的な性質のために、いくつかの側面でより優れたパフォーマンスを示しました。ランダムサーチはより迅速に検証器のバイアスに収束しましたが、他の2つのアルゴリズムでは最適ではない候補を改善する必要がありました。

ファインチューニングとの互換性

研究チームは、提案した検索方法とファインチューニングされたモデルとの互換性を検証しました。DPOファインチューニングされたStable Diffusion XLモデルを使用し、検索方法は異なるモデルに一般化でき、すでに調整されたモデルのパフォーマンスを向上させることができることがわかりました。

推論計算のさまざまな側面の影響

研究では、推論計算のさまざまな側面が結果にどのように影響するかを調査しました。

検索反復回数: 反復回数を増やすと、ノイズが最適値に近づきます。
検索反復ごとの計算: 反復ごとのノイズ除去ステップ数を調整すると、計算的に最適な領域が異なることが明らかになります。
最終生成計算: チームは、最終的なサンプル品質を最大化するために、最終的なノイズ除去ステップに最適な設定を使用しました。

計算投資の有効性

研究者らは、小規模な拡散モデルにおける推論時のスケーリングの有効性を調査しました。ImageNetの場合、小規模モデルのスケーリングは非常に効率的であることがわかりました。場合によっては、小規模モデルで検索する方が、検索なしで大規模モデルを使用するよりも優れたパフォーマンスを発揮する可能性があります。ただし、有効性は小規模モデルのベースラインパフォーマンスに依存します。

テキストベースの設定では、PixArt-Σは、FLUX-1.devよりも少ない計算量で優れたパフォーマンスを発揮しました。これらの結果は、トレーニング中に費やされた多大な計算リソースを、生成中の少量の計算で相殺できることを示しており、より効率的に高品質のサンプルが得られることを示しています。