基于分布层面奖励优化视觉生成模型

摘要

传统的视觉生成强化学习策略通常采用样本级奖励函数，但这一做法常导致奖励破解问题，从而降低图像多样性并引入视觉异常。为解决这些局限，我们提出了一种新框架，通过分布级奖励微调生成模型，确保其与真实数据分布更优对齐。与逐样本评估的奖励不同，分布级奖励考虑了样本的数据分布，有效缓解了所有样本独立朝相同方向优化时所出现的模式坍塌问题。为克服估计此类奖励的高昂计算成本，我们引入子集替换策略，通过仅更新生成参考集中的一小部分来高效提供奖励信号。此外，我们应用强化学习优化事后模型合并系数，从而可能缓解因在常规强化学习实践中引入随机微分方程（SDE）而导致的训练-推理不一致问题。大量实验表明，我们的方法显著提升了多种基模型的FID-50K指标：SiT从8.30降至5.77，EDM2从3.74降至3.52。定性评估也证实，该方法在保持样本多样性的同时增强了感知质量。

English

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.