V-GRPO: 노이즈 제거 생성 모델을 위한 온라인 강화학습은 생각보다 쉽다

기존의 이미지 생성 모델 강화학습은 복잡한 MDP(마르코프 결정 과정) 모델링으로 인해 학습이 매우 느리고 불안정했습니다. 이 논문은 단순한 ELBO 기반 대리 손실만으로도 적절한 안정화 기법을 더하면 기존 방식보다 2~3배 빠르고 우수한 성능을 낼 수 있음을 증명하여 생성 모델 사후 학습의 새로운 표준을 제시합니다.

기존의 Diffusion 모델 강화학습은 생성 과정을 여러 단계의 의사결정(MDP)으로 나누어 각 단계마다 보상을 최적화하려 했습니다. 이는 마치 긴 여정의 모든 발걸음을 일일이 교정하는 것과 같아 계산량이 방대하고 학습이 느려지는 원인이 되었습니다.

V-GRPO는 이 문제를 '전체 생성 결과'에 대한 확률(Log-likelihood)을 직접 최적화하는 방식으로 접근합니다. 이때 직접 계산이 불가능한 확률값 대신, Diffusion 모델의 학습 목표인 ELBO(증거 하한)를 대리 지표로 사용합니다. 이는 복잡한 경로 대신 최종 목적지의 점수를 기반으로 모델의 가중치를 한 번에 업데이트하는 것과 유사한 원리입니다.

하지만 ELBO는 샘플링 시점(Timestep)에 따라 값의 변동이 매우 커서 학습을 불안정하게 만듭니다. V-GRPO는 같은 프롬프트 그룹 내에서 노이즈를 공유하고, 모든 시점을 골고루 샘플링하며, 손실값의 크기를 스스로 조절하는 기법을 도입하여 이 불안정성을 제거했습니다. 결과적으로 기초적인 Gradient Descent 원리만으로도 고차원 이미지 생성 모델을 정교하게 튜닝할 수 있게 되었습니다.

V-GRPO는 GRPO의 Advantage 추정 방식과 Diffusion의 Variational Objective를 결합합니다. 먼저 behavior policy로부터 G개의 출력을 생성하고 각 출력에 대한 보상 Ri를 계산합니다. 이후 [Ri - mean(R)] / std(R) 연산을 통해 상대적 우위인 Advantage Ai를 구합니다. 이 Ai는 정책 업데이트 시 각 샘플의 영향력을 결정하는 가중치 역할을 수행합니다.

핵심 최적화는 중요도 샘플링 비율 rho를 계산할 때 발생합니다. [현재 모델의 ELBO - 이전 모델의 ELBO]를 지수 함수 exp()에 입력하여 rho 값을 얻습니다. 이 값은 현재 모델이 이전보다 해당 샘플을 생성할 확률이 얼마나 높아졌는지를 의미합니다. rho와 Ai를 곱한 값을 최대화하도록 모델을 학습시키되, rho가 너무 커지지 않도록 [1-epsilon, 1+epsilon] 범위로 clipping하여 급격한 변화를 방지합니다.

수치적 안정을 위해 Adaptive Loss Weighting을 적용합니다. [예측된 이미지와 실제 이미지의 차이(L2 norm)]를 [차이의 절댓값 평균(L1 norm)]으로 나누어 정규화합니다. 이 연산은 노이즈 수준에 따라 수십 배씩 차이 나는 손실값의 스케일을 일정하게 맞춰주어, 모든 학습 단계에서 균일한 gradient가 흐르도록 보장합니다.

FLUX.1-dev 모델 실험에서 V-GRPO는 300회 반복 학습만으로 HPS-v2.1(0.372), PickScore(0.241), ImageReward(1.749) 등 주요 지표에서 기존 최강자인 MixGRPO를 압도했습니다. 특히 MixGRPO가 25회의 샘플링 단계(NFE)를 요구할 때, V-GRPO는 단 4회의 NFE만으로도 더 높은 보상을 획득하며 효율성을 입증했습니다.

Stable Diffusion 3.5 Medium 환경에서도 DiffusionNFT가 1,700회의 업데이트를 거쳐 도달한 성능을 V-GRPO는 580회(약 3배 적은 횟수) 만에 달성했습니다. Ablation study 결과, 제안된 세 가지 분산 감소 기법 중 하나라도 누락될 경우 학습 곡선이 심하게 요동치거나 성능이 급격히 저하됨이 확인되어 각 구성 요소의 필수성이 증명되었습니다.

V-GRPO의 아키텍처는 별도의 Value Network 없이 그룹 내 상대적 보상을 사용하는 GRPO 구조를 따릅니다. 이는 파라미터 오버헤드를 줄이고 구현을 단순화합니다. 수학적으로는 Marginal Log-likelihood의 하한인 ELBO를 Surrogate로 채택하여, MDP 기반 방식이 가진 '전이 커널(Transition Kernel)에 대한 의존성'을 제거했습니다. 이 디커플링 덕분에 학습 시에는 1차 SDE 솔버를 쓰더라도 추론 시에는 고차원 ODE 솔버(DPM-Solver++ 등)를 자유롭게 선택할 수 있는 유연성을 확보했습니다.

구현 측면에서 Stratified Sampling은 전체 타임스텝 [0, 1] 구간을 NMC개의 동일 간격으로 나누고 각 구간에서 하나씩 샘플링하여 배치의 대표성을 높입니다. 또한 KL Penalty를 계산할 때 별도의 참조 모델 없이 이전 단계의 behavior policy를 활용하여 메모리 사용량을 최적화했습니다. Advantage Soft-clipping은 tanh 함수를 이용해 극단적인 보상값에 의한 gradient 폭주를 방지하며, 특히 단일 gradient step만 수행하는 완전 온폴리시(fully on-policy) 설정에서 학습 안정성을 크게 기여합니다.

논문은 V-GRPO가 GenEval과 같은 거친 입도(coarse-grained)의 보상 함수 환경에서는 기존의 중요도 샘플링 기반 방식보다 성능이 다소 낮을 수 있음을 언급했습니다. 또한 매우 적은 수의 타임스텝-노이즈 쌍(NMC < 4)을 사용할 경우 수렴 성능이 저하되는 포화 효과가 존재함을 명시했습니다.