Table of Contents
Fetching ...

GDRO: Group-level Reward Post-training Suitable for Diffusion Models

Yiyang Wang, Xi Chen, Xiaogang Xu, Yu Liu, Hengshuang Zhao

TL;DR

GDRO addresses the inefficiencies and instability of online RL in rectified-flow diffusion models by introducing a post-training, group-level reward optimization that is offline and sampler-independent. It derives an implicit reward function $s_\theta(x,t)$ from diffusion-DPO and uses a Plackett-Luce–inspired ranking loss to align a group of generated images with explicit rewards, augmented by a top-1 stabilization term. A corrected score is proposed to reflect reward hacking trends, enabling robust evaluation. Empirical results on OCR and GenEval show that GDRO improves reward scores efficiently while mitigating reward hacking, offering a practical path to reliable group-level reward alignment for diffusion models.

Abstract

Recent advancements adopt online reinforcement learning (RL) from LLMs to text-to-image rectified flow diffusion models for reward alignment. The use of group-level rewards successfully aligns the model with the targeted reward. However, it faces challenges including low efficiency, dependency on stochastic samplers, and reward hacking. The problem is that rectified flow models are fundamentally different from LLMs: 1) For efficiency, online image sampling takes much more time and dominates the time of training. 2) For stochasticity, rectified flow is deterministic once the initial noise is fixed. Aiming at these problems and inspired by the effects of group-level rewards from LLMs, we design Group-level Direct Reward Optimization (GDRO). GDRO is a new post-training paradigm for group-level reward alignment that combines the characteristics of rectified flow models. Through rigorous theoretical analysis, we point out that GDRO supports full offline training that saves the large time cost for image rollout sampling. Also, it is diffusion-sampler-independent, which eliminates the need for the ODE-to-SDE approximation to obtain stochasticity. We also empirically study the reward hacking trap that may mislead the evaluation, and involve this factor in the evaluation using a corrected score that not only considers the original evaluation reward but also the trend of reward hacking. Extensive experiments demonstrate that GDRO effectively and efficiently improves the reward score of the diffusion model through group-wise offline optimization across the OCR and GenEval tasks, while demonstrating strong stability and robustness in mitigating reward hacking.

GDRO: Group-level Reward Post-training Suitable for Diffusion Models

TL;DR

GDRO addresses the inefficiencies and instability of online RL in rectified-flow diffusion models by introducing a post-training, group-level reward optimization that is offline and sampler-independent. It derives an implicit reward function from diffusion-DPO and uses a Plackett-Luce–inspired ranking loss to align a group of generated images with explicit rewards, augmented by a top-1 stabilization term. A corrected score is proposed to reflect reward hacking trends, enabling robust evaluation. Empirical results on OCR and GenEval show that GDRO improves reward scores efficiently while mitigating reward hacking, offering a practical path to reliable group-level reward alignment for diffusion models.

Abstract

Recent advancements adopt online reinforcement learning (RL) from LLMs to text-to-image rectified flow diffusion models for reward alignment. The use of group-level rewards successfully aligns the model with the targeted reward. However, it faces challenges including low efficiency, dependency on stochastic samplers, and reward hacking. The problem is that rectified flow models are fundamentally different from LLMs: 1) For efficiency, online image sampling takes much more time and dominates the time of training. 2) For stochasticity, rectified flow is deterministic once the initial noise is fixed. Aiming at these problems and inspired by the effects of group-level rewards from LLMs, we design Group-level Direct Reward Optimization (GDRO). GDRO is a new post-training paradigm for group-level reward alignment that combines the characteristics of rectified flow models. Through rigorous theoretical analysis, we point out that GDRO supports full offline training that saves the large time cost for image rollout sampling. Also, it is diffusion-sampler-independent, which eliminates the need for the ODE-to-SDE approximation to obtain stochasticity. We also empirically study the reward hacking trap that may mislead the evaluation, and involve this factor in the evaluation using a corrected score that not only considers the original evaluation reward but also the trend of reward hacking. Extensive experiments demonstrate that GDRO effectively and efficiently improves the reward score of the diffusion model through group-wise offline optimization across the OCR and GenEval tasks, while demonstrating strong stability and robustness in mitigating reward hacking.
Paper Structure (19 sections, 19 equations, 12 figures, 3 tables)

This paper contains 19 sections, 19 equations, 12 figures, 3 tables.

Figures (12)

  • Figure 1: Illustrations of our method. As shown in the top block, GDRO optimizes the text-to-image model toward high precision OCR rendering and layout planning, while maintaining high quality and prompt-image alignment. In contrast, in the bottom block, although achieving a high evaluation score, Flow-GRPO hacks the reward to make the text very big and the content of the image unnatural.
  • Figure 2: Overview of our method. Given a pre-generated image group synthesized from the same prompt and their corresponding explicit rewards, we perturb the images with noise on different time steps, feed them to the diffusion model to predict the velocity, and calculate the implicit rewards accordingly to get the final loss.
  • Figure 3: Reward misalignment and hacking. The first row shows the reward misalignment scenario. The second row shows two reward hacking cases.
  • Figure 4: Qualitative results on OCR and GenEval. Columns 1-4 show the OCR task, and Columns 5-8 display the GenEval task.
  • Figure 5: Evaluation curves. We plot the evaluation time scores and corrected scores of different methods on the OCR and GenEval task across GPU training hours. We mark the reward hacking dividing line of Flow-GRPO on its peak corrected scores.
  • ...and 7 more figures