通过分布级奖励优化视觉生成模型

Optimizing Visual Generative Models via Distribution-wise Rewards

精选理由

这篇论文提出用分布级奖励替代样本级奖励,解决了生成模型奖励作弊的老问题,SiT和EDM2的FID分数都显著提升。

AI 摘要

论文提出分布级奖励(Distribution-wise Rewards)框架,替代传统样本级奖励函数,以解决视觉生成中的奖励作弊和模式坍塌问题。通过子集替换策略降低计算成本,并应用强化学习优化模型合并系数,缓解训练-推理不一致。在SiT模型上FID-50K从8.30降至5.77,在EDM2模型上从3.74降至3.52。定性评估显示该方法在保持多样性的同时提升了感知质量。

原文 · arXiv cs.LG

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.