GARDO: Reinforcing Diffusion Models without Reward Hacking

TL;DR

GARDO tackles reward hacking in RL-based fine-tuning of diffusion models by introducing selective regularization, adaptive anchoring, and diversity-aware optimization. It gates KL regularization to samples with high reward uncertainty, updates the reference policy adaptively to keep the regularization target relevant, and reshapes advantages to reward high-quality, diverse samples, promoting broader mode coverage. Empirical results across proxy and unseen metrics show GARDO mitigates hacking while preserving sample efficiency and exploration, achieving or surpassing baseline performance and improving diversity. The approach demonstrates robust improvements across multiple proxy rewards and unseen tasks, indicating practical impact for safer, more reliable text-to-image alignment. The method relies on auxiliary reward models for uncertainty estimation and faces open questions for scaling to video generation and other modalities.

Abstract

Fine-tuning diffusion models via online reinforcement learning (RL) has shown great potential for enhancing text-to-image alignment. However, since precisely specifying a ground-truth objective for visual tasks remains challenging, the models are often optimized using a proxy reward that only partially captures the true goal. This mismatch often leads to reward hacking, where proxy scores increase while real image quality deteriorates and generation diversity collapses. While common solutions add regularization against the reference policy to prevent reward hacking, they compromise sample efficiency and impede the exploration of novel, high-reward regions, as the reference policy is usually sub-optimal. To address the competing demands of sample efficiency, effective exploration, and mitigation of reward hacking, we propose Gated and Adaptive Regularization with Diversity-aware Optimization (GARDO), a versatile framework compatible with various RL algorithms. Our key insight is that regularization need not be applied universally; instead, it is highly effective to selectively penalize a subset of samples that exhibit high uncertainty. To address the exploration challenge, GARDO introduces an adaptive regularization mechanism wherein the reference model is periodically updated to match the capabilities of the online policy, ensuring a relevant regularization target. To address the mode collapse issue in RL, GARDO amplifies the rewards for high-quality samples that also exhibit high diversity, encouraging mode coverage without destabilizing the optimization process. Extensive experiments across diverse proxy rewards and hold-out unseen metrics consistently show that GARDO mitigates reward hacking and enhances generation diversity without sacrificing sample efficiency or exploration, highlighting its effectiveness and robustness.

Figures (14)

• Figure 1: Setting OCR as the proxy reward, vanilla RL methods like Flow-GRPO exploit the OCR reward at the cost of losing real image quality, leading to reward hacking. It generates unrealistic, noisy images with blurry backgrounds and visual artifacts. In contrast, our method maintains better image quality and diversity. Prompt:"A storefront with 'GARDO' written on it".
• Figure 2: Overview of GARDO. GARDO introduces an uncertainty-driven, gated KL mechanism to control the proportion of regularization, avoiding unnecessary penalties. Our proposed diversity-aware advantage shaping effectively encourages exploration of novel states.
• Figure 3: We train a diffusion model with $3$-layer MLP on Gaussian mixtures (pre-trained distribution), with the goal to capture multimodal high-reward clusters as shown in the reward landscape. The vanilla RL method (DDPO DDPO) with a large KL coefficient $\beta$ is overly constrained and fails to increase rewards. Conversely, a small $\beta$ incurs severe mode collapse. Our proposed diversity-aware optimization, when applied alone, successfully captures the multimodal modes, including the central cluster with the lowest probability density in the reference policy $\pi_{\rm ref}$. Our full GARDO framework simultaneously achieves maximum reward and discovers all high-reward clusters.
• Figure 4: Learning curves and o.o.d. generalization results across different methods. GARDO not only matches the sample efficiency of the KL-free baseline, but also mitigates reward hacking effectively, as evidenced by the superior performance on unseen metrics.
• Figure 5: Our diversity-aware advantage shaping effectively improves the generation diversity.

Theorems & Definitions (2)

• Definition 1
• Proposition 1