[論文レビュー] Diffusion Reinforcement Learning via Centered Reward Distillation

Yuanzhi Zhu, Xi Wang+2 more|arXiv (Cornell University)|Mar 14, 2026

Reinforcement Learning in Robotics被引用数 0

ひとこと要約

CRD は拡散 RL を中心化された報酬マッチング問題として定式化し、未知のプロンプト依存正規化項を打ち消し、外部報酬を用いたテキストから画像へのファインチューニングをより競争力があり、より速く、より堅牢に実現する安定化技法を導入する。

ABSTRACT

Diffusion and flow models achieve State-Of-The-Art (SOTA) generative performance, yet many practically important behaviors such as fine-grained prompt fidelity, compositional correctness, and text rendering are weakly specified by score or flow matching pretraining objectives. Reinforcement Learning (RL) fine-tuning with external, black-box rewards is a natural remedy, but diffusion RL is often brittle. Trajectory-based methods incur high memory cost and high-variance gradient estimates; forward-process approaches converge faster but can suffer from distribution drift, and hence reward hacking. In this work, we present \textbf{Centered Reward Distillation (CRD)}, a diffusion RL framework derived from KL-regularized reward maximization built on forward-process-based fine-tuning. The key insight is that the intractable normalizing constant cancels under \emph{within-prompt centering}, yielding a well-posed reward-matching objective. To enable reliable text-to-image fine-tuning, we introduce techniques that explicitly control distribution drift: (\textit{i}) decoupling the sampler from the moving reference to prevent ratio-signal collapse, (\textit{ii}) KL anchoring to a CFG-guided pretrained model to control long-run drift and align with the inference-time semantics of the pre-trained model, and (\textit{iii}) reward-adaptive KL strength to accelerate early learning under large KL regularization while reducing late-stage exploitation of reward-model loopholes. Experiments on text-to-image post-training with \texttt{GenEval} and \texttt{OCR} rewards show that CRD achieves competitive SOTA reward optimization results with fast convergence and reduced reward hacking, as validated on unseen preference metrics.

研究の動機と目的

• Diffusion モデルを静的な事前学習目的を超えた外部報酬へポスト訓練で整列させることを動機付ける。
• 解くのが困難な正規化項にもかかわらず、順方向プロセス拡散 RL の頑健な目的関数を開発する。
• RL ファインチューニング中の報酬ハックと分布ドリフトを低減する。
• 非分化可能な報酬を用いた安定的で効率的な拡散 RL を実現する実用的な安定化技法を提供する。

提案手法

• Centered Reward Distillation (CRD) を KL 正規化報酬最大化の下で中心化された報酬マッチング目的として導出する。
• 現在のモデルと移動参照間の暗黙の対数密度比を推定するために前方拡散 ELBO の代理を用いる。
• 各プロンプトグループ内で報酬を中心化してプロンプト依存の正規化項（beta log Z(c)）を打ち消す。
• 報酬に依存する重みと温度パラメータ tau によって報酬加重中心化重みを定義する。
• 比率信号の崩壊を防ぐためにサンプリングを移動参照から切り離す。
• CFG ガイド付き事前訓練モデルへの KL アンカーリングを適用して長期的な漂移を抑制し、推論時の意味論と整合させる。
• 報酬適応型 KL 強度を導入して早期学習を加速しつつ後期の報酬の抜け穴を緩和する。

実験結果

リサーチクエスチョン

• RQ1 未知の計算困難な正規化項を知らなくても、定義が適切な拡散 RL の目的をどのように形成できるか。
• RQ2 プロンプト内での中心化がプロンプト依存の正規化を除去し、報酬マッチングを安定化させるか。
• RQ3 分布ドリフトと報酬ハックに対して前方プロセス拡散 RL を最も安定化させる実践的技法は何か。
• RQ4 CRD は Flow-GRPO 型および前方プロセスベースのベースラインと比べて、競争力のある報酬最適化とより速い収束、改善された学習安定性を達成するか。
• RQ5 CFG ベースの KL アンカーリングは訓練の安定性と推論時の意味論の整合性にどのように影響するか。

主な発見

• CRD は報酬をプロンプトグループ内で中心化することにより打ち消し不可能な正規化項を除去し、適切な報酬マッチング目的を得る。
• CRD はベースラインと比較して競争力のある報酬最適化とより速い収束・改善された訓練安定性を達成する。
• 実用的な安定化技法—分離サンプリング、CFG ガイド付き事前訓練モデルへの KL アンカーリング、報酬適応型 KL 強度—は報酬ハックと分布ドリフトを緩和する。
• GenEval および OCR 報酬を用いたテキストから画像へのファインチューニングの実験は、質的・量的に有利な挙動と報酬ハックの低減を示す。
• CRD は従来の報酬蒸留および GVPO 風の目的を特別なケースとして包含し、InfoNCA に関連する比率ベースの変種を許容する。