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.

한 줄 요약

시각 생성 모델의 분포 기반 보상 최적화를 통해 FID-50K 3.74 → 3.52 개선.

핵심 기여도

샘플 기반 보상이 모드 콜랩스와 시각 이상을 유발함을 분석.
분포 기반 보상과 subset-replace 전략을 도입하여 FID-50K 3.74 → 3.52 개선.
SDE 기반 학습과 ODE 기반 추론 간 불일치를 post-hoc 모델 병합 최적화로 해결.
SiT 모델에서 FID-50K 8.30 → 5.77, FD-DINOv2 230.39 → 164.88 개선.

핵심 아이디어

기존 시각 생성 모델의 강화학습은 샘플 기반 보상 함수를 사용하지만, 이는 보상 해킹을 유발해 이미지 다양성 저하와 시각 이상을 초래한다. 이를 해결하기 위해, 본 연구는 생성 샘플의 전체 분포를 고려하는 **분포 기반 보상**(distribution-wise reward)을 도입한다. 이는 단일 샘플이 아닌 전체 샘플 분포와 목표 분포 간의 차이를 평가하여 **모드 콜랩스**(mode collapse)를 완화한다.

구체적으로, **subset-replace 전략**을 통해 기존 참조 집합의 일부를 새로 생성된 샘플로 교체하고, 교체된 집합의 FID를 보상 신호로 사용한다. 이는 대규모 샘플 집합을 필요로 하지 않아 계산 비용을 절감한다. 또한, SDE 기반 학습과 ODE 기반 추론 간 불일치를 해결하기 위해, **post-hoc 모델 병합 계수**(model merging coefficients)를 강화학습을 통해 최적화한다.

기술적 접근법

**분포 기반 보상**: FID(Fréchet Inception Distance)를 사용하여 생성 샘플 분포와 목표 분포 간 차이를 평가.
**subset-replace 전략**: 참조 집합의 0.01× 크기의 작은 부분 집합을 새로 생성된 샘플로 교체하고, 교체된 집합의 FID를 보상으로 사용.
**post-hoc 모델 병합 최적화**: 강화학습을 통해 모델 병합 계수를 최적화하여 학습-추론 불일치를 해결.
**모델**: SiT, EDM2 사용.
**하이퍼파라미터**: 참조 집합은 클래스 균형을 고려한 중간 크기의 생성 샘플로 구성.

주요 결과

**SiT 모델**: FID-50K 8.30 → 5.77 (+29.9%), FD-DINOv2 230.39 → 164.88 (+28.2%).
**EDM2 모델**: FID-50K 3.74 → 3.52 (+5.9%), EDM2-XS 2.57 → 2.52 (+1.9%).
**시각 평가**: 생성 이미지의 품질 향상과 다양성 유지가 확인됨.

의의 및 한계

본 연구는 샘플 기반 보상의 한계를 극복하고, 생성 모델의 분포적 일관성과 다양성을 동시에 향상시키는 새로운 강화학습 프레임워크를 제시한다. 특히, **subset-replace 전략**은 계산 효율성을 확보하면서도 강력한 보상 신호를 제공하며, **post-hoc 모델 병합 최적화**는 학습-추론 불일치를 해결하는 실용적 접근법이다.

그러나, 본 연구는 주로 이미지 생성 분야에 초점을 맞추고 있으며, 동일한 접근법이 다른 생성 태스크(예: 영상, 텍스트)에 적용 가능한지는 추가 연구가 필요하다. 또한, FID는 일부 경우 실제 인간 판단과의 괴리가 존재할 수 있으므로, 보다 정교한 분포 평가 지표의 도입이 필요할 수 있다.

실용적 활용

본 연구는 생성 모델의 학습-추론 불일치를 해결하고, 생성 이미지의 품질과 다양성을 동시에 향상시키는 데 기여할 수 있다. 특히, **EDM2와 SiT 같은 대형 생성 모델**의 post-hoc 최적화에 유용하며, **디지털 콘텐츠 생성, 인공지능 아트, 광고 이미지 생성** 등 다양한 산업 분야에서 활용 가능하다.