透過分佈層面獎勵優化視覺生成模型

摘要

傳統的視覺生成強化學習策略通常採用樣本層級的獎勵函數，但此做法經常引發獎勵駭客行為，導致影像多樣性降低並產生視覺異常。為解決這些限制，我們提出一個新穎框架，透過基於分佈的獎勵來微調生成模型，確保與真實資料分佈更佳對齊。與逐個樣本評估的獎勵不同，基於分佈的獎勵考量樣本的資料分佈，從而緩解當所有樣本獨立朝相同方向最佳化時所發生的模式崩塌問題。為克服估算此類獎勵所需的龐大計算成本，我們引入子集替換策略，僅透過更新生成參考集中的一小部分子集，即可有效提供獎勵訊號。此外，我們應用強化學習來最佳化事後模型合併係數，可能緩解因在常規強化學習實作中引入隨機微分方程（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.