Table of Contents
Fetching ...

Data-regularized Reinforcement Learning for Diffusion Models at Scale

Haotian Ye, Kaiwen Zheng, Jiashu Xu, Puheng Li, Huayu Chen, Jiaqi Han, Sheng Liu, Qinsheng Zhang, Hanzi Mao, Zekun Hao, Prithvijit Chattopadhyay, Dinghao Yang, Liang Feng, Maosheng Liao, Junjie Bai, Ming-Yu Liu, James Zou, Stefano Ermon

TL;DR

This work tackles reward hacking in reinforcement learning for diffusion models by identifying on-policy regularization as a core weakness. It introduces Data-regularized Diffusion Reinforcement Learning (DDRL), a forward KL-based framework that anchors learning to an off-policy data distribution and directly links to diffusion training, enabling seamless post-training integration. Through large-scale video and image experiments, DDRL consistently improves task rewards while aligning with human preferences, and it remains effective with synthetic data. The approach provides a principled, scalable foundation for diffusion RL at scale, with theoretical guarantees and practical advantages in data efficiency and post-training versatility.

Abstract

Aligning generative diffusion models with human preferences via reinforcement learning (RL) is critical yet challenging. Most existing algorithms are often vulnerable to reward hacking, such as quality degradation, over-stylization, or reduced diversity. Our analysis demonstrates that this can be attributed to the inherent limitations of their regularization, which provides unreliable penalties. We introduce Data-regularized Diffusion Reinforcement Learning (DDRL), a novel framework that uses the forward KL divergence to anchor the policy to an off-policy data distribution. Theoretically, DDRL enables robust, unbiased integration of RL with standard diffusion training. Empirically, this translates into a simple yet effective algorithm that combines reward maximization with diffusion loss minimization. With over a million GPU hours of experiments and ten thousand double-blind human evaluations, we demonstrate on high-resolution video generation tasks that DDRL significantly improves rewards while alleviating the reward hacking seen in baselines, achieving the highest human preference and establishing a robust and scalable paradigm for diffusion post-training.

Data-regularized Reinforcement Learning for Diffusion Models at Scale

TL;DR

This work tackles reward hacking in reinforcement learning for diffusion models by identifying on-policy regularization as a core weakness. It introduces Data-regularized Diffusion Reinforcement Learning (DDRL), a forward KL-based framework that anchors learning to an off-policy data distribution and directly links to diffusion training, enabling seamless post-training integration. Through large-scale video and image experiments, DDRL consistently improves task rewards while aligning with human preferences, and it remains effective with synthetic data. The approach provides a principled, scalable foundation for diffusion RL at scale, with theoretical guarantees and practical advantages in data efficiency and post-training versatility.

Abstract

Aligning generative diffusion models with human preferences via reinforcement learning (RL) is critical yet challenging. Most existing algorithms are often vulnerable to reward hacking, such as quality degradation, over-stylization, or reduced diversity. Our analysis demonstrates that this can be attributed to the inherent limitations of their regularization, which provides unreliable penalties. We introduce Data-regularized Diffusion Reinforcement Learning (DDRL), a novel framework that uses the forward KL divergence to anchor the policy to an off-policy data distribution. Theoretically, DDRL enables robust, unbiased integration of RL with standard diffusion training. Empirically, this translates into a simple yet effective algorithm that combines reward maximization with diffusion loss minimization. With over a million GPU hours of experiments and ten thousand double-blind human evaluations, we demonstrate on high-resolution video generation tasks that DDRL significantly improves rewards while alleviating the reward hacking seen in baselines, achieving the highest human preference and establishing a robust and scalable paradigm for diffusion post-training.

Paper Structure

This paper contains 46 sections, 1 theorem, 23 equations, 8 figures, 6 tables.

Key Result

Theorem 3.1

Maximizing eq:ddrl_objective_diff is equivalent to maximizing eq:ddrl_objective, and its optimal policy $p_\theta^*(\mathbf{x})$ satisfies

Figures (8)

  • Figure 1: Generated videos after RL with different algorithms from the same checkpoint. DDRL improves rewards by generating prompt-aligned and realistic videos, while DanceGRPO and FlowGRPO increase rewards with over-colorized, over-stylized, and unrealistic videos.
  • Figure 2: (Left) Reward hacking during training. All methods are trained from the same base model using the same dataset for 256 iterations. While DanceGRPO and FlowGRPO achieve higher rewards, (Right) humans consistently prefer videos generated by the base model. By contrast, DDRL improves rewards and its generations are more likely preferred than the base model's generations.
  • Figure 3: Even with significantly larger $\beta$ and the KL divergence remains stable throughout the training, unrealistic noise textures still appear in videos generated by the latter checkpoint.
  • Figure 4: Breakdown of the the increase/decline of text alignment (TA), visual quality (VQ), and motion quality (MQ), after post-training the 2B model with different methods. DanceGRPO shows a 16% (T2V) and 28% (I2V) decrease in the TA score, despite its average reward is the highest. This aligns with the human voting result where its generated videos are not preferred. DDRL is the only algorithm that brings "Pareto improvement".
  • Figure 5: VideoAlign reward after post-training with DDRL from (1) pretrained model and (2) SFT model. The former achieves comparably high rewards with significantly higher data efficiency.
  • ...and 3 more figures

Theorems & Definitions (1)

  • Theorem 3.1