Fine-Tuning Diffusion Models via Intermediate Distribution Shaping

この論文は、拡散モデルの中間ノイズレベルにおける分布を操作する手法（P-GRAFT）と、明示的な報酬なしで学習誤差を修正する「逆ノイズ補正」アルゴリズムを提案し、これらが既存の手法やベースモデルよりもテキストから画像生成や分子生成などのタスクにおいて優れた性能を示すことを実証しています。

要約

この論文は、**「AI が絵を描くとき、どうすればもっと上手に、そして効率的に描けるようになるか」**という問いに答える、画期的な新しい方法を提案しています。

AI が絵を描く技術（拡散モデル）は、最初は「真っ白なノイズ（砂嵐のようなもの）」から始めて、少しずつノイズを取り除いていき、最終的に美しい絵を完成させます。このプロセスを「絵の具を混ぜていく作業」や「彫刻家が石から像を彫り出す作業」に例えるとわかりやすいかもしれません。

この論文では、その「作業の途中」に新しい工夫を加えることで、AI の能力を劇的に向上させる方法を 2 つ提案しています。

---

1. 「途中の形」を整える魔法：P-GRAFT

（Partial-GRAFT：部分的重み付けサンプリング）

🎨 従来の方法の課題

これまでの AI の学習では、「完成した絵」を見て、「これはいい絵だ（報酬が高い）」と判断し、その結果を元に AI を修正していました。
しかし、完成した絵を見るのは、**「石を完全に彫り終えてから、初めて『これは像だ』と判断する」**ようなものです。石を削りながらの途中経過（ノイズの多い状態）では、それが最終的にどんな絵になるか予測が難しく、AI が「どう直せばいいか」を学ぶのが大変で、エラーも起きやすくなります。

✨ 新しい方法：途中の形に注目する

この論文が提案するP-GRAFTは、「完成した絵」ではなく、「途中の半分の形」に注目して学習させるという考え方です。

• アナロジー：
  Imagine you are teaching a student to draw a cat.

  • 旧来の方法: 学生が完成した絵を描くまで待ってから、「猫の耳が変だ」と指摘する。
  • P-GRAFT の方法: 学生が「まだ輪郭しか描いていない状態（途中）」で、「この形なら、最終的に猫になりそうだ！」と判断し、その途中の形を褒めて学習させる。

• なぜ素晴らしいのか？
  途中の形（ノイズが少し取れた状態）は、最終的な完成形に比べて**「予測が簡単」で「誤差（バイアス）」が小さい**からです。
  論文ではこれを「バイアスとバリアンスのトレードオフ（偏りと揺らぎのバランス）」と呼んでいます。

  • 完成形に近いほど、何が正解か（報酬）はわかりやすいが、AI がそれを学ぶのは難しい（複雑すぎる）。
  • 途中の形なら、AI が学ぶのは簡単だが、何が正解か（報酬）は少し曖昧。
  • P-GRAFTは、この「学びやすさ」と「正解の明確さ」の**絶妙なバランス点（途中のステップ）**を見つけることで、AI がより効率的に、より上手に絵を描けるようにします。

結果： テキストから画像を作るテスト（Stable Diffusion など）で、従来の方法よりもはるかに高い精度で、指示された通りの絵を描けるようになりました。

---

2. 「逆さまのノイズ」を直す：Inverse Noise Correction

（流体力学モデル向けの新しいアプローチ）

🌊 従来の方法の課題

絵を描く AI には「拡散モデル」の他に「フローモデル」という種類もあります。これは、「最初にあるノイズ（スタート地点）」を決めれば、その後の絵は自動的に決まるという仕組みです。
しかし、AI が学習する過程で、「スタート地点のノイズ」の性質が少しずれてしまうことがあります。

• アナロジー：
  川を下って海（完成した絵）にたどり着く船があるとします。
  本来は「広大な海（標準的なノイズ）」から出発するはずが、学習の過程で「狭い川（歪んだノイズ）」から出発してしまい、結果として海にたどり着いた船の位置が少しずれてしまいます。

✨ 新しい方法：スタート地点を修正する

この論文が提案するInverse Noise Correction（逆ノイズ補正）は、「完成した絵」から逆算して、「正しいスタート地点（ノイズ）」を特定し、それを補正するという方法です。

• 仕組み：

  1. 完成した絵（データ）を、AI に逆方向に流して「元々のノイズ」に戻します。
  2. その「戻ってきたノイズ」を見て、「あ、ここがずれているな」と判断します。
  3. 新しい小さな AI（ノイズ補正器）を作って、「正しいノイズ」を「ずれたノイズ」に変換するように学習させます。
  4. 本物の絵を描くときは、まずこの補正器でノイズを直し、それから本物の AI に絵を描かせます。

• メリット：

  • 特別な評価基準が不要： 通常、AI を良くするには「いい絵・悪い絵」を人間が評価する必要がありますが、この方法は**「評価なし」**で、AI 自体の仕組みを直すだけで品質を上げられます。
  • 計算コストの削減： 完成するまで長い時間をかける必要がなくなり、「より少ない計算量（FLOPs）」で、より高画質な絵が描けるようになります。

---

まとめ：この研究がもたらすもの

この論文は、AI 絵画の「学習の仕方」を根本から変える 2 つのアイデアを提示しました。

1. P-GRAFT（途中の形を整える）：
   「完成してから直す」のではなく、「途中の段階で方向を修正する」ことで、AI がより賢く、指示通りに描けるようにする。

   • 例： 彫刻家が石を削りながら、完成形をイメージして微調整する。

2. Inverse Noise Correction（スタート地点を直す）：
   「完成した結果」から逆算して、「出発点のノイズ」を修正することで、評価基準がなくても AI の質を上げ、計算コストも下げる。

   • 例： 地図が少しずれていることに気づき、出発地点を補正して目的地に正確に着くようにする。

これらの技術により、AI はより少ない計算資源で、より高品質な画像を生成できるようになり、医療、デザイン、科学など、さまざまな分野での活用がさらに加速することが期待されています。

技術概要

論文「FINE-TUNING DIFFUSION MODELS VIA INTERMEDIATE DISTRIBUTION SHAPING」の技術的サマリー

この論文は、事前学習済み拡散モデル（Diffusion Models）やフローモデル（Flow Models）の微調整（Fine-tuning）において、最終的なデータ分布だけでなく、拡散過程の中間ノイズレベルにおける分布を整形（Shaping）することの重要性を提唱し、そのための新しいフレームワークとアルゴリズムを提案しています。

以下に、問題定義、手法、主要な貢献、実験結果、および意義について詳細にまとめます。

---

1. 背景と問題定義

背景

拡散モデルは生成タスクで広く利用されていますが、事前学習済みモデルを特定のタスク（例：プロンプトへの整合性向上、エラー修正）に合わせて微調整する必要があるケースが増えています。従来のアプローチでは、強化学習（RL）アルゴリズム、特に KL 正則化付き報酬最大化（PPO など）が用いられてきました。

既存手法の課題

1. KL 正則化の非計算可能性: 拡散モデルでは、マルギナル尤度（Marginal Likelihood）が計算不可能（intractable）であるため、PPO における KL 正則化項を正確に計算できません。
2. 既存の近似手法の限界:
   • KL 項を無視すると大規模設定で学習が不安定になる。
   • 軌道 KL（Trajectory KL）を緩和条件として用いると、初期値関数のバイアス問題が発生し、結果が劣化する。
3. 拒否サンプリング（Rejection Sampling）の未活用: 拒否サンプリングに基づく微調整手法（RAFT など）は存在するが、これらが KL 正則化された報酬最大化とどのように関連するか、また拡散モデルの中間状態をどう活用するかが十分に解明されていなかった。

---

2. 提案手法

著者らは、拒否サンプリングと KL 正則化報酬最大化の間の概念的なつながりを明らかにし、2 つの主要な手法を提案しています。

2.1 GRAFT (Generalized Rejection sAmpling Fine-Tuning)

まず、既存の拒否サンプリング手法（RAFT, Best-of-N など）を「一般化拒否サンプリング（GRS）」という共通フレームワークで統合しました。

• 理論的発見: GRS を用いてサンプルを選択・微調整することは、再整形された報酬（Reshaped Rewards）を用いた KL 正則化報酬最大化を暗黙的に実行することに等しいことを数学的に示しました。
• 意義: これにより、マルギナル尤度が計算不可能な拡散モデルに対しても、KL 制約を満たす微調整が可能になります。

2.2 P-GRAFT (Partial-GRAFT): 中間分布整形

拡散モデルの特性（時間経過とともにノイズが除去される過程）を利用し、最終生成物だけでなく、中間ステップ（Partial Denoising）における分布を整形する手法を提案しました。

• 仕組み:
  1. 事前学習モデルから完全な生成サンプル（$X_0$）を生成し、報酬を計算します。
  2. 高報酬のサンプルを選択しますが、微調整対象とするのは**中間時点 $t$ のノイズ状態（$X_t$）**です。
  3. 選択された $X_t$ を用いて、$t$ から $T$（ノイズ開始点）までの部分のみを微調整します。
  4. 推論時には、微調整済みモデルで $T \to t$ のデノイジングを行い、その後は参照モデル（Reference Model）で $t \to 0$ を行います。
• 理論的根拠（バイアス - バラつきトレードオフ）:
  • バラつき（Variance）: 中間ステップ $t$ における報酬の条件付き分散は、$t$ が大きくなる（ノイズに近い）ほど増大します（報酬ノイズが増える）。
  • バイアス（Bias）: 一方で、$t$ が大きい領域（ノイズに近い領域）では、スコア関数（Score Function）が単純なガウス分布のスコアに近づくため、学習が容易になります（バイアス減少）。
  • 結論: 適切な中間ステップ $t$ を選ぶことで、学習の難易度（バイアス）と報酬のノイズ（バラつき）のバランスを最適化でき、より効果的な微調整が可能になります。

2.3 逆ノイズ補正 (Inverse Noise Correction)

フローモデル（Flow Models）向けに、明示的な報酬なしで事前学習モデルの学習誤差を修正する手法を提案しました。

• 仕組み:
  1. 事前学習済みフローモデルを用いて、データ分布から逆方向（データ $\to$ ノイズ）に ODE を逆向きに解き、「逆ノイズ分布（Inverse Noise Distribution）」を推定します。
  2. この逆ノイズ分布を学習する新しいアダプターモデル（Noise Corrector）を訓練します。
  3. 推論時には、標準ガウスノイズをまず Noise Corrector で「逆ノイズ」に変換し、その後、元の事前学習モデルで画像を生成します。
• 意義: 事前学習モデルの最終分布の誤差は、初期ノイズ分布の誤差と等価であるという性質を利用し、ノイズ分布を補正することで生成品質を向上させます。

---

3. 主要な貢献

1. GRAFT フレームワークの確立: 拒否サンプリングに基づく微調整が、KL 正則化報酬最大化の解を再整形された報酬で実現することを証明し、拡散モデルへの適用を可能にしました。
2. P-GRAFT の提案と理論的正当化: 中間分布を整形する新しい手法 P-GRAFT を提案し、バイアス - バラつきトレードオフの観点からその有効性を数学的に説明しました。
3. 逆ノイズ補正の導入: 報酬なしでフローモデルの誤差を修正するパラメータ効率の良い手法を開発しました。
4. 広範な実験評価: テキストから画像（T2I）、レイアウト生成、分子生成、無条件画像生成など、多様なタスクで手法を検証しました。

---

4. 実験結果

テキストから画像生成 (Text-to-Image)

• モデル: Stable Diffusion v2 (SDv2) をベースに微調整。
• 評価指標: VQAScore（プロンプトと画像の整合性）。
• 結果:
  • P-GRAFT は、PPO ベースの手法（DDPO）やベースモデル（SDv2, SDXL-Base）を大幅に上回りました。
  • 特に $NI = 0.25N$（中間ステップを早期に設定）の構成が最も性能が高く、VQAScore でベースモデルに対し8.81% の相対改善を達成しました。
  • DDPO は大規模なプロンプトセットでは学習が不安定になり性能が低下しましたが、P-GRAFT は安定して高品質な結果を出しました。

レイアウト生成と分子生成

• レイアウト (PubLayNet): P-GRAFT は無条件・条件付き生成の両方で、FID スコアとアライメント指標を改善しました。
• 分子生成 (QM9): 安定した分子の生成率を向上させました。特に、多様性を保つための「デデュプリケーション（重複除去）」を組み合わせた P-GRAFT が、9 倍のサンプル数を使った GRAFT よりも少ない計算コストで高い安定性を示しました。

無条件画像生成 (逆ノイズ補正)

• データセット: CelebA-HQ, LSUN-Church。
• 結果: 逆ノイズ補正を適用することで、FID スコアが大幅に改善されました。
• 効率性: 計算量（FLOPs）を削減しつつ、より高い品質を達成しました。例えば、ノイズ補正モデル（100 ステップ）＋事前学習モデル（100 ステップ）の組み合わせは、事前学習モデル単体（1000 ステップ）よりも優れた FID を示しました。

---

5. 意義と結論

この論文は、拡散モデルの微調整において「最終結果」だけでなく「生成過程の中間状態」を制御することの重要性を理論的・実証的に示しました。

• 理論的貢献: 拒否サンプリングと KL 正則化 RL の等価性を示し、拡散モデル特有の中間分布整形によるバイアス - バラつきトレードオフを明らかにしました。
• 実用的貢献: 既存の強化学習手法（PPO など）よりも安定して高性能な微調整手法（P-GRAFT）を提供し、また報酬なしでモデルを改善する新しいアプローチ（逆ノイズ補正）を提案しました。
• 応用範囲: 連続的な画像生成だけでなく、離散・連続混合モデル（分子・レイアウト）やフローモデルにも汎用的に適用可能です。

総じて、この研究は生成モデルの微調整における新しいパラダイム（中間分布の制御）を確立し、計算効率と生成品質の両面で大きな進歩をもたらすものです。