HKU GDRO Fixes Diffusion Reward Cheating

🔑 Enhanced Key Takeaways

• •GDRO introduces a 'Corrected Score' metric that explicitly penalizes 'Metric Collapse,' a phenomenon where models achieve high OCR/GenEval scores by generating oversized text or simplified backgrounds that sacrifice overall image detail.
• •The method is 'Sampler-Independent,' meaning it does not require the conversion of deterministic ODEs to stochastic SDEs; this avoids the out-of-domain artifacts and quality degradation common in online RL methods like Flow-GRPO.
• •GDRO utilizes 'Implicit Reward Functions' computable at any diffusion timestep, enabling gradient updates without a separate differentiable reward model or a critic network, which significantly reduces memory overhead.
• •The framework shifts the training paradigm from online sampling to an offline 'Trajectory-Reward' approach, allowing the model to learn from pre-computed denoising paths rather than requiring real-time image generation during the RL loop.

🛠️ Technical Deep Dive

• •Group-level Normalization: Implements a minimax risk principle across a group of N samples for a single prompt, estimating advantages using normalized relative rewards without a value function.
• •Implicit Reward Manipulation: Proves theoretically that reward optimization can be performed by manipulating implicit functions derived from the diffusion score at any timestep (t).
• •Offline Importance Sampling: Uses importance sampling to correct for distribution shifts between the behavior policy (dataset generation) and the target policy, enabling stable offline learning.
• •Corrected Evaluation Logic: The evaluation framework integrates a 'hacking trend' coefficient that adjusts the final score based on the divergence from the original model's visual fidelity.
• •Architecture Compatibility: Specifically optimized for the Vision Transformer (ViT) backbone of FLUX.1-dev, targeting the 96% of compute spent in the transformer module.

🔮 Future ImplicationsAI analysis grounded in cited sources

⏳ Timeline