분포별 보상을 통한 시각 생성 모델 최적화

초록

기존의 시각적 생성을 위한 강화학습 전략은 일반적으로 표본 단위 보상 함수를 사용하지만, 이러한 접근법은 종종 보상 해킹을 초래하여 이미지 다양성을 저하시키고 시각적 이상을 발생시킨다. 이러한 한계를 해결하기 위해, 우리는 분포 단위 보상을 사용하여 생성 모델을 미세 조정함으로써 실제 데이터 분포와의 정렬을 개선하는 새로운 프레임워크를 제시한다. 표본을 개별적으로 평가하는 보상과 달리, 분포 단위 보상은 표본들의 데이터 분포를 고려하여 모든 표본이 독립적으로 동일한 방향으로 최적화될 때 발생하는 모드 붕괴 문제를 완화한다. 이러한 보상을 추정하는 데 따르는 막대한 계산 비용을 극복하기 위해, 우리는 생성된 참조 집합의 작은 부분 집합만을 업데이트하여 효율적으로 보상 신호를 제공하는 부분 집합 교체 전략을 도입한다. 또한, RL을 적용하여 사후 모델 병합 계수를 최적화함으로써, 일반적인 RL 관행에서 확률적 미분 방정식(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.