Optimizing Visual Generative Models via Distribution-wise Rewards

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.

AI image generators learn to create realistic images, but improving them with feedback is tricky. Current approaches rate each generated image individually and reward the model for producing higher-scoring outputs. However, this causes the model to "game" the scoring system. It learns to exploit shortcuts that inflate scores without truly improving quality, leading to repetitive or distorted images. We propose a fundamentally different approach: instead of judging images one by one, we evaluate whether a collection of generated images, taken together, resembles real photographs as a whole. This encourages the model to produce diverse, realistic outputs rather than repeatedly exploiting the same trick. To make this computationally practical, we developed a strategy that swaps only a small batch of images into a larger reference collection and measures how the overall quality changes. Our method reduces a standard quality metric (FID) by up to 30% on multiple models, producing images that are both more realistic and more varied without needing extra training data or architectural changes.