技術的サマリー：Reinforce Adjoint Matching (RAM)

問題定義

拡散モデルおよびフローマッチングモデルは、その事前学習が教師あり回帰に依存しているため、高忠実度な生成を実現します。具体的には、クリーンなサンプルを解析的にノイズ化し、モデルが閉形式のターゲットに対して回帰を行います。しかし、これらのモデルを報酬関数（例：構成的正しさ、テキスト描画、人間の好意など）に合わせて、強化学習（RL）による事後学習を行うと、この効率的な構造が崩れてしまいます。

既存の RL 事後学習手法は、計算上および構造的な重大な障壁に直面しています：

1. 方策勾配法: 去ノイズをマルコフ決定過程として扱い、確率的な SDE ロールアウトを必要とし、変分下限を最適化します。これにより高い分散が生じ、多数のサンプリングステップが必要となります。
2. 確率的最適制御（随伴）法: KL 正則化付きの報酬最大化目的を正確にターゲットとしますが、保存された SDE 軌道に沿った後向き随伴掃引と、サンプリング過程を通じて逆伝播される報酬勾配を必要とします。これは数値的に不安定で計算コストが高く、特にノイズスケジュールがノイズの多い終端付近で発散するフローマッチングモデルにおいて顕著です。
3. 代理損失: 最近のアプローチは、扱いにくい尤度を代理（ELBO や対照的エンドポイントなど）に置き換えますが、これにより事前学習のスケーラビリティを可能にするクリーンな回帰構造が犠牲になります。

核心的な課題は、事前学習の回帰フレームワークの単純性とスケーラビリティを維持しつつ、報酬の最適化を効果的に行う RL 事後学習目的を導出することです。

手法：Reinforce Adjoint Matching (RAM)

著者らは、SDE ロールアウト、後向き随伴掃引、または報酬勾配を必要とせずに、事前学習の回帰構造を RL 事後学習に拡張する手法**Reinforce Adjoint Matching (RAM)**を提案します。

理論的基盤

この手法は、確率的最適制御問題として定式化されたKL 正則化付き報酬最大化に基づいています。

1. 最適プロセスの構造: KL 正則化の下では、最適生成プロセスはクリーンなエンドポイント ($X_0$) の分布を高報酬サンプル ($p^*_0 \propto p_{ref}^0 \exp(r)$) 方向に傾けますが、クリーンなエンドポイントが与えられたときのノイズ状態の条件付き法則は変化しません。具体的には、固定された $X_0$ からノイズ状態 $X_t$ への軌道は、事前学習で使用された解析的なガウスカーネルと同一のままです。
2. 随伴一致条件: 最適制御 $u_t$ は、制御が価値関数（随伴）の負のスケール付き勾配に等しいという固定点条件を満たします：$u_t = -\sigma_t \nabla_x V_t(x)$。
3. REINFORCE 分解: $\nabla r$（微分不可能な場合がある）の計算を避けるため、著者らは対数微分（REINFORCE）恒等式を価値関数に適用します。これにより、随伴は、後向きブリッジスコアで重み付けされた報酬項と、パスコスト補正項に分解されます。

RAM 目的関数

RAM は、パスコスト補正を除外することで随伴を近似します（これは低分散と大規模での安定性によって正当化されます）。また、ノイズ化カーネルが解析的であるという構造的洞察を利用します。

1. エンドポイントサンプリング: 制御された SDE をシミュレートする代わりに、RAM は既存の ODE サンプリャーを使用して、現在のモデルの分布からクリーンなエンドポイント $X_0$ をサンプリングします。
2. 解析的ノイズ化: サンプリングされた $X_0$ を解析的にノイズ化して、事前学習と同様にトレーニング状態 $X_t = (1-t)X_0 + t\epsilon$ を生成します。
3. 閉形式ターゲット: 速度場 $v_\theta$ の回帰ターゲットは、報酬項とベイズブリッジスコア（線形ノイズ化に対して閉形式式を持つ）を随伴条件に代入することで導出されます。

得られる損失関数は以下の通りです：
$$L_{RAM}(\theta) = \mathbb{E}_t \left[ \left\| v_\theta^t(X_t) - \text{sg}\left( v_{ref}^t(X_t) + r(X_0) \left( (\epsilon - X_0) - v_\theta^t(X_t) \right) \right) \right\|^2 \right]$$
ここで、$X_0 \sim p_\theta^0$、$\epsilon \sim \mathcal{N}(0, I)$、$\text{sg}(\cdot)$ はストップグラディエント演算子を表します。

主要な効率化メカニズム

• SDE ロールアウトの不要: トレーニング状態は、独立したノイズベクトルを用いて単一のエンドポイントを $K$ 回ノイズ化することで生成されます。これにより、モデルの前方パスあたり $K$ 個の条件付き独立トレーニングサンプルが得られ、単一の SDE 軌道に依存する手法よりもステップあたりの勾配信号が豊富になります。
• 報酬勾配の不要: この手法はブラックボックス報酬で機能します。
• 後向き掃引の不要: ノイズ化カーネルの解析的性質により、後向き随伴積分の必要性が排除されます。

主要な貢献

1. 構造的拡張: 本論文は、KL 正則化の下では事前学習の回帰構造が RL 事後学習に拡張されることを実証しています。最適プロセスはエンドポイント分布のみを変更し、条件付きノイズ化法則は変更しません。
2. RAM アルゴリズム: 閉形式のベイズブリッジスコアを用いて報酬情報で事前学習ターゲットを修正する新規の一貫性損失であり、既存手法の計算ボトルネックを回避します。
3. 理論的正当性: RAM の固定点が KL 正則化最適解に対応することを示す導出であり、線形報酬下でのガウス参照に対しては完全性が保証され、一般的には一次一致が保証されます。
4. 分散分析: パスコスト推定量の分析により、パスコスト補正を除外すること（RAM が行うこと）が、完全範囲または局所随伴推定量と比較して分散を大幅に減少させることが示され、画像スケールの事後学習において実用的な唯一の推定量となっています。

実験結果

著者らは、RAM をStable Diffusion 3.5 Medium (SD3.5M) において以下の 3 つのタスクで評価しました：

1. 構成的画像生成 (GenEval): 物体、属性、空間関係の正しさを測定。
2. 視覚的テキスト描画 (OCR): 生成画像内のテキストの可読性を測定。
3. 人間の好意への整合 (PickScore): 人間の好意への整合性を測定。

パフォーマンス:

• 報酬最大化: RAM は、Flow-GRPO、DiffusionNFT、AWM を含むベースラインを上回り、すべてのタスクで最高トレーニング報酬を達成しました。
• トレーニング効率: RAM は、GenEval で Flow-GRPO のピーク報酬に到達するのに50 倍少ないステップ、OCR で48 倍少ないステップ、PickScore で34 倍少ないステップで達成しました。
• 画像品質: タスク指標のために美的品質を低下させる「報酬ハッキング」に苦しんだ一部のベースラインとは異なり、RAM はタスク報酬を最大化しながら画像品質指標（Aesthetic、DeQA、ImageReward、HPSv2）を維持または向上させました。
• CFG 互換性: RAM は事後学習後も Classifier-Free Guidance (CFG) と互換性を保ちますが、DiffusionNFT や AWM などのベースラインは CFG を適用するとタスク報酬が大幅に低下しました。

意義と主張

本論文は、RAM が拡散モデルおよびフローマッチングモデルにおける RL 事後学習のパラダイムシフトを表すと主張しています。最適制御問題の特定の構造的性質を活用することで、RAM は事前学習をスケーラブルにした「単純な回帰」パラダイムを回復させます。

著者らは以下を主張しています：

• RAM は、高解像度画像生成への RL スケーリングにおける主要なボトルネックであった高コストな SDE ロールアウトと後向き随伴掃引を排除します。
• この手法は「単純でスケーラブル」であり、既存の手法と比較してトレーニングステップを大幅に削減しつつ、ステップあたりの計算コストは同等で、最先端の結果を達成します。
• このアプローチは、可視的な品質劣化や報酬ハッキングなしに、複雑な報酬（構成的正しさ、テキスト、好意）とモデルを整合させることに成功しており、正しい理論的条件が満たされれば、事前学習構造が事後学習の整合性を処理するのに十分に頑健であることを実証しています。