技術概要：Diffused Reward を用いた Diff-Instruct（Didr）

問題定義

最近のワンステップ型テキストから画像（T2I）生成モデルの進展により、潜在ノイズを単一のフォワードパスで直接画像にマッピングすることでリアルタイム合成が可能となった。しかし、これらの生成モデルを強化学習（RL）を通じて人間の好みに整合させることは、根本的な構造的課題を伴う。

従来のワンステップ型 RLHF の手法（例：Diff-Instruct*、Diff-Instruct++）は、画像空間における報酬最適化と潜在空間における分布一致を組み合わせている。これらのアプローチは通常、クリーンな画像の終点（$x_0$）でのみ選好信号（報酬）を適用し、拡散過程全体のノイズ軌道にわたって参照分布に対するカルバック・ライブラー（KL）正則化を課す。著者らは、このパラダイムに決定的な構造的ミスマッチが存在することを特定した。すなわち、報酬は終端状態でのみ作用する一方、正則化はすべてのノイズレベルにまたがって適用されるという点である。

このミスマッチは、「終端報酬支配（terminal reward domination）」と呼ばれる失敗モードを引き起こす。前方ノイズは自然に画像の詳細を隠蔽するため、高ノイズレベルではモードの発散に対する KL 正則化ペナルティが著しく弱まる。その結果、最適化アルゴリズムはこの軌道の脆弱性を悪用し、十分な KL ペナルティを負うことなく参照分布を放棄して高報酬領域へ収束してしまう。これにより報酬スコアは向上するが、しばしば画像の忠実度と分布の安定性が犠牲となる。

手法：Diffused Reward を用いた Diff-Instruct（Didr）

この不整合を解決するため、著者らは積分 KL（IKL）最小化から導出されたデータフリーの軌道レベル整合フレームワークであるDidrを提案する。核心的な洞察は、RLHF 最適の報酬傾斜付きクリーン画像分布を、拡散軌道のすべてのノイズレベルにわたって体系的に伝播させることである。

1. 原理的な軌道ターゲット

この手法は、KL 正則化付き RLHF におけるクリーン画像の正しい整合ターゲットを確立することから始まる。最適クリーン画像分布 $q^*(x_0|c)$ は、報酬傾斜付き密度として以下のように表される。
$$q^*(x_0|c) \propto q_0(x_0|c) \exp(r(x_0, c)/\tau)$$
ここで、$q_0$ は参照分布、$r$ は報酬関数、$\tau$ は温度パラメータである。

報酬を $x_0$ のみで適用するのではなく、Didr はこの最適ターゲットを参照フォワード過程を通じて $q^*$ を拡散させることで、完全な軌道に「持ち上げる」。これにより、すべてのノイズレベルで定義された報酬傾斜付き周辺分布 $\{q^*_t\}_{t \in [0, T]}$ の族が得られる。
$$q^*_t(x_t|c) = \int q_t(x_t|x_0) q^*(x_0|c) dx_0$$
著者らは、生成器の周辺分布とこれらの報酬傾斜付き周辺分布との間の積分 KL 発散を最小化することが、標準的な KL 正則化付き RLHF 目的関数と同じ最小化器を持つことを証明した。これにより、報酬と正則化がすべてのノイズレベルでバランスされ、終端報酬支配が防止される。

2. Diffused Reward Score（DRS）

IKL 目的関数を微分することで、著者らはターゲットスコア関数を導出した。最適スコアは 2 つの成分に分解される。
$$s_{target} = s_{ref} + s_r$$
ここで、$s_{ref}$ は参照スコア関数、$s_r$ は**Diffused Reward Score（DRS）**である。DRS はすべてのノイズレベルで定義される報酬駆動型の補正項として機能する。
$$s_r(x_t, t, c) = \nabla_{x_t} \log \mathbb{E}_{x_0 \sim q(x_0|x_t, c)} \left[ \exp\left(\frac{r(x_0, c)}{\tau}\right) \right]$$
分類器ガイダンス（一次近似）とは異なり、DRS は完全な対数指数傾斜を保持し、ノイズ観測からクリーン画像が識別可能な場合のみ選好情報を伝播させる。

3. Diffused Reward Proxy（DRP）

DRS は扱いにくい事後期待値 $q(x_0|x_t, c)$ を必要とするため、著者らは手法を実用的にするために**Diffused Reward Proxy（DRP）**を導入した。DRP は、凍結された参照モデルからの微分可能な短ステップ復元連鎖を用いて事後期待値を近似する。

• メカニズム: ノイズ状態 $x_t$ から、凍結された参照モデルを用いて $K$ 個の独立した $S$ ステップ復元連鎖を実行し、近似クリーンサンプル $\{\hat{x}^{(k)}_0\}$ を生成する。
• 勾配推定: DRS は、パスごとの報酬勾配のソフトマックス重み付き平均として近似される。
  $$\hat{s}_r(x_t, t, c) = \frac{1}{\tau} \sum_{k=1}^K \omega^{(k)} \nabla_{x_t} r(\hat{x}^{(k)}_0, c)$$
  ここで、重み $\omega^{(k)}$ は $\exp(r(\hat{x}^{(k)}_0, c)/\tau)$ に比例する。
• 効率性: このアプローチにより、参照モデルの活性化を保存したり画像トレーニングデータを必要としたりすることなく、報酬勾配が軌道を通じて安定して逆伝播することを可能にする。

4. トレーニングフレームワーク

トレーニングプロセスは 2 つの段階を交互に繰り返す。

1. Teaching Assistant（TA）更新: TA スコアモデル（$s_\psi$）を、生成器の進化中の周辺分布を追跡するようにデノイジング・スコアマッチング（DSM）で訓練し、暗黙的な生成器スコアの代理とする。
2. 生成器更新: 生成器（$g_\theta$）を、TA スコアと CFG 修正済み参照スコアおよび DRP 推定量の合計を比較する IKL 目的関数の勾配を用いて更新する。

主要な貢献

1. 終端報酬支配の特定: 本論文は、最適化アルゴリズムがノイズレベルを悪用して高報酬領域へ収束し、意図された RLHF のバランスを侵害する、従来のワンステップ型 RLHF 手法における構造的失敗モードを形式的に特定・分析した。
2. 原理的な軌道レベル整合: 著者らは、RLHF 最適のクリーン画像密度を前方へ伝播させることで、数学的に妥当な軌道レベルターゲットを導出するフレームワーク Didr を提案した。これにより、報酬と正則化がすべてのノイズレベルでバランスされることが保証される。
3. DRS の解析的導出: 本論文は、Diffused Reward Score を解析的に導出し、軌道ターゲットが参照スコアと報酬誘発型補正に正確に分解されることを示した。
4. 実用的推定量（DRP）: 著者らは、微分可能な短ステップ復元を通じて DRS を計算するゼロデータ推定量である Diffused Reward Proxy を導入し、画像データなしで安定したトレーニングを可能にした。

実証結果

広範な実験により、Didr が既存のワンステップ型ベースラインを一貫して上回ることが示された。

• パレート支配: SDXL バックボーン（1024×1024）において、標準的な Didr 変種は、PickScore–FID トレードオフにおいて既存のすべてのワンステップ型整合ベースラインをパレート支配し、FLUX-dev や SDXL-DPO などのマルチステップモデルを両軸で同時に凌駕した。
• 選好指標: 拡張トレーニング変種（Didr_longer）は、評価されたワンステップモデルの中で最高の PickScore（23.9）を達成し、最も強力な先行ベースライン（Diff-Instruct*）に対して ImageReward を 8.9% 改善した。また、HPSv2.1 および Aesthetic Score においても首位を走った。
• 汎化性: Didr は 60 億パラメータの Z-Image バックボーンへも成功裡に転送された。単一の生成ステップにおいて、選好指標でその 50 ステップの教師モデルを上回り、アーキテクチャ全体にわたるスケーラビリティを実証した。
• 定性的改善: 視覚的比較により、Didr は従来のワンステップ型手法に比べてより鮮明な微細なディテールと自然な色調を生成し、報酬最適化されたベースラインでよく見られる過剰飽和や絵画的なアーティファクトを回避することが示された。

意義と主張

本論文は、原理的な軌道レベル報酬伝播が、エンドポイントのみの報酬手法に内在する不整合を解決する効果的な整合戦略であると主張している。RLHF 最適分布を拡散過程を通じて厳密に伝播させることで、Didr はワンステップ生成器の中で選好整合と画像忠実度の間の最良のトレードオフを達成する。

著者らは、この手法が画像トレーニングデータを一切必要とせず、テキストプロンプトと凍結された参照モデルのみに依存することを強調している。彼らは Didr を、ワンステップ生成器がマルチステップの教師モデルと同等かそれ以上の選好整合を達成しつつ、リアルタイム推論速度を維持することを可能にする解決策として位置づけている。この研究は、先行手法で特定された構造的ミスマッチが、終端状態だけでなく軌道全体を整合させることでのみ解決可能な根本的な制限であることを示唆している。