技術要約：報酬を用いたモーメント・マッチング蒸留による拡散モデルのファインチューニング

問題提起

拡散モデルは、その安定した学習目的関数により、高精度な画像合成や科学的予測における標準となっています。しかし、実用的なデプロイメントにおいては、数十から数百のデノイジング・ステップを必要とする推論の計算コストの高さが障壁となっています。**蒸留（Distillation）**技術はこれらのモデルを数ステップの生成器へと圧縮し、強化学習（RL）によるファインチューニングは、特定の報酬関数（人間の好みや科学的指標など）にモデルを適合させますが、これら2つのフェーズを組み合わせることは依然として困難です。

蒸留と報酬最適化を統合しようとする既存のアプローチには、以下のような重大な限界があります：

1. 結合学習における脆弱性: 蒸留と報酬ファインチューニングを単一のフェーズに統合すると、生成サンプルが教師モデルの入力分布から逸脱し、蒸留の信号が無効化されることがよくあります。
2. メモリとバイアスの問題: 多段階のバックプロパゲーション（例：DRaFT）を用いて既に蒸留されたモデルをファインチューニングすることは、メモリ消費が激しくなります。一方、バックプロパゲーションを打ち切る手法（例：ReFL）は、ノイズの乗った潜在変数に対して報酬を評価するため、バイアスが生じます。
3. 構造的制約: HyperNoiseのような手法は、初期ノイズを摂動させることでバックプロパゲーションの連鎖を回避しますが、低周波の画像変化に限定されており、複雑な高周波の構造的報酬を最適化する能力には限界があります。

手法：報酬を用いたモーメント・マッチング蒸馏 (RMMD)

著者らは、蒸留と報酬ファインチューニングを原理的に結びつけ、高度な蒸留が持つ高精度な「自然さ」を維持しつつ、RMMD (Rewarded Moment Matching Distillation) と呼ばれる2フェーズのフレームワークを提案しています。

フェーズ1：モーメント・マッチング蒸留 (MMD)

プロセスは、ベースとなる教師モデルをモーメント・マッチング蒸留 (MMD) を用いて蒸留することから始まります。この手法は、サンプリング・軌跡に沿った中間デノイジング分布を一致させることで、教師の周辺分布を密接に追跡するマルチステップの生徒モデルを生み出します。この生徒モデルは、安定した分布的参照として機能するために固定されます。

フェフェーズ2：オンポリシー報酬ファインチューニング

第2フェーズでは、分布の忠実性を維持しながら、報酬関数を最大化するように生徒モデルをファインチューニングします。主な革新点は以下の通りです：

1. オンポリシー・サンプリング: オフポリシーのデータポイントを使用する代わりに、現在の生徒のポリシーからノイズの乗った状態 $x^{pol}_t$ をサンプリングします（生徒を $K_{student}$ ステップ実行し、出力を再ノイズ化し、勾配を停止することで実現）。これにより、報酬目的関数が生徒の期待報酬を直接近似することを保証し、分布のドリフトによって学習信号が無効化されるのを防ぎます。
2. リサイクルされた正則化: コアとなる貢献は、モーメント・マッチング損失をファインチューニング中の正則化項として再利用することです。これは単なるヒューリスティックなペナルティではなく、生徒の中間デノイジング分布と固定された参照モデルとの間の不一致を明示的に罰するものです。これにより、報酬の最大化と分布の忠実性の間のトレードオフを制御できる解釈可能なノブ（正則化の重み）を提供します。
3. シングルステップ・グラディエント: 破損したオンポリシー・サンプルに対して単一のデノイジング・ステップを行い、報酬を評価することで、多段階の展開を行うことなく全周波数帯域を捉えます。
4. L2 正則化: 初期 MMD 蒸留モデルの予測からの大幅な逸脱を罰するために、追加の L2 項を導入し、学習をさらに安定させます。

総たる目的関数は、オンポリシー報酬勾配と、モーメント・マッチング正則化、および L2 項を組み合わせたものです：
$$\nabla_\theta \mathcal{L}_{RMMD} = \mathbb{E} \left[ \left(\frac{\partial \hat{x}_0}{\partial \theta}\right)^\top \left( m_{student} - m_{teacher} - \lambda \nabla_{\hat{x}_0} R(\hat{x}_0) \right) \right] + \lambda_{reg} \mathcal{L}_{L2}$$

主な貢献

1. RMMD 手順: 分布のドリフトを防ぐためのオンポリシー版モーメント・マッチング損失を利用した、新しい拡散蒸留手法。
2. 経験的な優位性: ピクセルレベルの指標からセマンティックな整合性（CLIP）に至る多様な報酬関数において、RMMD が FID-Reward パレートフロント上で、シングルステップのベースライン（DI++）やマルチステップの競合手法（DRaFT, HyperNoise）を一貫して上回ることを実証しました。
3. 科学的応用: RMMD を最先端の拡散ベース気象予測モデルである GenCast に適用し、連続ランク確率スコア（CRPS）を最適化することに成功しました。

結果

ImageNet 実験

• マルチステップの優位性: 8ステップの MMD 蒸留生徒を起点とすることで、1ステップのベースライン（DI++）を大幅に上回りました。8ステップの RMMD は、優れた FID-Reward パレートフロントを実現しました。
• 比較: RMMD は、様々な解像度（64x64 および 512x512）および報酬タイプ（CLIP 整合性、Inception スコア、滑らかさ）において、DRaFT や HyperNoise よりも優れたパレートフロントを生成しました。
• 定性的分析: 可視化の結果、DRaFT-1 はコンテンツを保持するものの敵対的なアーティファクトを導入し、DRaFT-2 は報酬ハッキング（劇的な分布シフト）に苦しむ一方で、RMMD は元のデータ分布から大きく逸脱することなく、報酬を高めるための微細な修正を統合することに成功しました。

GenCast 気象予測

• 高速化: 蒸留されたモデルは、教師モデルと比較して 7.5倍の高速化（8ステップ vs 59 回の関数評価）を達成しました。
• 精度: RMMD で最適化されたモデルは、ターゲットとなる気象変数の 93% において教師モデルを上回りました（MMD のみの最良ベースラインは 75%）。
• 較正（キャリブレーション）: RMMD は、MMD 蒸留で一般的な過小分散の問題に対処しました。オンポリシー版はオフポリシー版よりもわずかに劣る分散指標を示しましたが、全体的な CRPS の改善においては最良の結果を提供しました。
• 汎化性能: 12時間リードタイムでの CRPS 最適化は、最大7日間の予測においても一貫した改善をもたらしました。これは、本手法が単に直近の報酬に過学習しているのではなく、基礎となる分布モデリングを改善していることを示唆しています。

意義と主張

本論文は、RMMD が、結合された蒸留とファインチューニングに固有の「ドリフト」問題に対する原理的な解決策を提供すると主張しています。モーメント・マッチング損失を正則化として再利用し、オンポリシー・サンプリングを利用することで、生成の質や蒸留モデルの安定性を損なうことなく、複雑で微分可能な目的関数（気象予測における CRPS など）の最適化が可能になります。

著者らは、このアプローチが高次元の複雑な科学領域にもスケールすることに重点を置いています。GenCast への適用成功は、拡散モデルが確率的な予測をより効率的かつ正確に改善できることを示しており、微分可能な目的関数が存在する科学計算における、より効率的で正確な拡散ベースモデルへの道を開くものです。本研究は、モーメント・マッチングを通じて安定した分布的参照を維持することが、拡散モデルにおける効果的な報酬ファインチューニングにとって極めて重要であることを示唆しています。