分布単位の報酬による視覚生成モデルの最適化

要旨

従来の視覚生成における強化学習手法は、通常サンプル単位の報酬関数を採用するが、この手法はしばしば報酬ハッキングを引き起こし、画像の多様性を低下させ、視覚的な異常をもたらす。これらの限界に対処するため、本稿では分布に基づく報酬を用いて生成モデルを微調整する新しいフレームワークを提案する。これにより、実世界のデータ分布とのより良い整合性が確保される。サンプルを個別に評価する報酬とは異なり、分布に基づく報酬はサンプルのデータ分布を考慮し、すべてのサンプルが独立に同じ方向へ最適化されることで生じるモード崩壊の問題を軽減する。これらの報酬の推定に伴う莫大な計算コストを克服するため、生成された参照セットの小さなサブセットのみを更新することで効率的に報酬信号を提供するサブセット置換戦略を導入する。さらに、事後的なモデルマージング係数を最適化するために強化学習を適用し、通常の強化学習手法における確率微分方程式（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.