Optimizing Visual Generative Models via Distribution-wise Rewards

2026-07-02

AI Summary

This work addresses the issue of reward hacking in traditional reinforcement learning, where sample-level rewards often lead to mode collapse and visual artifacts. To mitigate this, the authors propose a distribution-level reward mechanism that enhances generation quality by aligning the generated distribution with the true data distribution. They introduce an efficient subset replacement strategy to approximate the reward, significantly reducing computational overhead. Furthermore, reinforcement learning is employed to optimize posterior model fusion coefficients, alleviating the train-inference inconsistency problem. The method yields substantial improvements in FID-50K across multiple baseline models—e.g., reducing SiT from 8.30 to 5.77 and EDM2 from 3.74 to 3.52—while preserving sample diversity and enhancing perceptual quality.

Abstract

Conventional reinforcement learning strategies for visual generation typically employ sample-wise reward functions, yet this practice frequently results in reward hacking that degrades image diversity and introduces visual anomalies. To address these limitations, we present a novel framework that finetunes generative models using distribution-wise rewards, ensuring better alignment with real-world data distributions. Unlike rewards that evaluate samples individually, distribution-wise reward accounts for the data distribution of the samples, mitigating the mode collapse problem that occurs when all samples optimize towards the same direction independently. To overcome the prohibitive computational cost of estimating these rewards, we introduce a subset-replace strategy that efficiently provides reward signals by updating only a small subset of a generated reference set. Additionally, we apply RL to optimize post-hoc model merging coefficients, potentially mitigating the train-inference inconsistency caused by introducing stochastic differential equation (SDE) in regular RL practices. Extensive experiments show our approach significantly improves FID-50K across various base models, from 8.30 to 5.77 for SiT and from 3.74 to 3.52 for EDM2. Qualitative evaluation also confirms that our method enhances perceptual quality while preserving sample diversity.

Problem

• reward hacking
• mode collapse
• sample diversity
• train-inference inconsistency
• visual anomalies

Innovation

• distribution-wise rewards
• reward hacking
• mode collapse
• subset-replace strategy
• model merging