Diffusion Reinforcement Learning via Centered Reward Distillation
Computer Vision and Pattern RecognitionArtificial IntelligenceMachine Learning
Summary
The authors study how to improve the quality of images generated from text by using a method called reinforcement learning (RL) to fine-tune diffusion models. They find that existing RL methods can be unstable or cheat the reward system, so they propose a new approach called Centered Reward Distillation (CRD) that carefully controls how the model learns from rewards. Their method helps the model learn faster and avoid tricks that ruin image quality. Tests show that CRD achieves strong results with less training trouble and better alignment with human preferences.
Authors
Yuanzhi Zhu, Xi Wang, Stéphane Lathuilière, Vicky Kalogeiton
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.