확산 모델을 위한 강화 학습 설계 공간 재고: 손실 함수 설계 외의 가능도 추정의 중요성에 대한 연구

Rethinking the Design Space of Reinforcement Learning for Diffusion Models: On the Importance of Likelihood Estimation Beyond Loss Design

강화 학습은 텍스트-이미지 생성과 같은 시각적 작업에서 확산 및 흐름 모델에 널리 적용되어 왔습니다. 그러나 확산 모델은 계산하기 어려운 가능도를 가지고 있어, 기존의 정책 경사 방법론을 직접 적용하는 데 어려움이 있습니다. 기존 연구는 주로 이미 정교하게 설계된 LLM 기반 목표 함수를 활용하고, 가능도 추정을 위해 임의적인 방법을 사용하며, 이러한 추정이 전체 알고리즘 성능에 미치는 영향을 충분히 조사하지 않았습니다. 본 연구에서는 정책 경사 목표 함수, 가능도 추정 방법, 그리고 시뮬레이션 샘플링 방식의 세 가지 요소를 분리하여 강화 학습 설계 공간을 체계적으로 분석합니다. 연구 결과, 최종 생성 샘플에서만 계산되는 ELBO(Evidence Lower Bound) 기반 모델 가능도 추정기가 효과적이고 효율적이며 안정적인 강화 학습 최적화를 가능하게 하는 가장 중요한 요소이며, 특정 정책 경사 손실 함수의 영향보다 더 크다는 것을 확인했습니다. SD 3.5 Medium을 사용하여 여러 보상 벤치마크에서 연구 결과를 검증했으며, 모든 작업에서 일관된 경향을 관찰했습니다. 제안하는 방법은 GenEval 점수를 0.24에서 0.95로 향상시켰으며, 이는 90시간의 GPU 시간 동안 달성된 결과입니다. 이는 FlowGRPO보다 4.6배 효율적이고, SOTA(State-of-the-Art) 방법인 DiffusionNFT보다 2배 효율적이며, 보상 조작 없이 달성된 결과입니다.

Reinforcement learning has been widely applied to diffusion and flow models for visual tasks such as text-to-image generation. However, these tasks remain challenging because diffusion models have intractable likelihoods, which creates a barrier for directly applying popular policy-gradient type methods. Existing approaches primarily focus on crafting new objectives built on already heavily engineered LLM objectives, using ad hoc estimators for likelihood, without a thorough investigation into how such estimation affects overall algorithmic performance. In this work, we provide a systematic analysis of the RL design space by disentangling three factors: i) policy-gradient objectives, ii) likelihood estimators, and iii) rollout sampling schemes. We show that adopting an evidence lower bound (ELBO) based model likelihood estimator, computed only from the final generated sample, is the dominant factor enabling effective, efficient, and stable RL optimization, outweighing the impact of the specific policy-gradient loss functional. We validate our findings across multiple reward benchmarks using SD 3.5 Medium, and observe consistent trends across all tasks. Our method improves the GenEval score from 0.24 to 0.95 in 90 GPU hours, which is $4.6\times$ more efficient than FlowGRPO and $2\times$ more efficient than the SOTA method DiffusionNFT without reward hacking.