UDM-GRPO: 균등 이산 확산 모델을 위한 안정적이고 효율적인 그룹 상대 정책 최적화

이산 확산 모델(UDM)에 강화학습을 직접 적용할 때 발생하는 학습 불안정성과 성능 저하 문제를 해결한 최초의 프레임워크입니다. 정확한 액션 정의와 궤적 재구성을 통해 텍스트-이미지 생성 및 OCR 성능을 획기적으로 개선하여 이산 도메인 생성 모델의 새로운 학습 표준을 제시합니다.

기존의 확산 모델은 노이즈가 섞인 상태에서 점진적으로 노이즈를 제거하며 이미지를 생성한다. 하지만 이 과정에 강화학습을 적용하면, 초기 단계의 부정확한 예측값(Action)을 정답처럼 흉내 내야 하는 모순이 발생한다. 특히 이산 확산 모델(UDM)에서는 토큰 단위의 변화가 급격하여 학습이 쉽게 붕괴(Collapse)되는 한계가 있었다.

UDM-GRPO는 이 문제를 해결하기 위해 '결과 중심의 학습' 방식을 채택한다. 모델이 매 순간 내놓는 중간 결과물이 아니라, 그 결과물로부터 예측한 '최종 완성본'을 액션으로 간주한다. 이를 통해 모델은 항상 보상이 높은 최종 목표를 향해 가중치를 업데이트하게 된다.

또한, 학습 시 모델이 직접 생성한 불안정한 경로 대신, 원본 이미지에 노이즈를 섞어 만든 안정적인 경로(Forward Process)를 따라가도록 설계했다. 이는 모델이 사전 학습 때 익혔던 데이터 분포를 벗어나지 않게 잡아주는 가이드라인 역할을 수행하며, 결과적으로 훨씬 빠르고 안정적인 수렴을 가능하게 한다.

UDM-GRPO는 T-step MDP를 재구성하여 정책 최적화를 수행한다. 핵심은 액션 a_t를 최종 클린 샘플 x_1의 예측값으로 정의하고, 정책 π_θ(a_t | s_t)를 현재 상태 x_t에서 x_1이 나올 확률 p_θ(x_1 | x_t)로 설정한 것이다.

[현재 노이즈 상태 x_t 입력] → [모델이 최종 이미지 x_1의 각 토큰별 확률(Logits) 출력] → [Softmax 연산 수행] → [해당 타임스텝의 정책 확률 π_θ 결정]. 이 방식은 Euler solver의 효율성을 유지하면서도 미분 가능성을 확보하여 그라디언트 전파를 원활하게 한다.

학습 궤적은 X_forward 방식을 사용한다. [최종 샘플 x_1에 노이즈를 추가] → [중간 상태 x_t 생성] → [이 경로를 학습 데이터로 사용] → [사전 학습 분포와의 FID 거리 최소화]. 이는 모델이 생성한 궤적(X_backward)이 훈련 데이터 분포에서 벗어나는 OOD(Out-of-Distribution) 문제를 방지한다.

URSA 모델에 적용한 결과, GenEval 벤치마크에서 종합 점수가 0.69에서 0.96으로 크게 상승하며 SOTA를 달성했다. 특히 객체 수 세기(Counting) 성능은 0.60에서 0.95로, 속성 결합(Attr. Binding)은 0.49에서 0.85로 비약적인 향상을 보였다.

OCR 성능 지표에서도 기존 8% 수준의 정확도를 57%까지 끌어올려 텍스트 렌더링 능력을 입증했다. 인간 선호도 지표인 PickScore 역시 20.46에서 23.81로 증가했다. Ablation study를 통해 Accurate Action과 Forward Trajectory 전략이 없을 경우 학습 500단계 이후 KL 발산이 폭발하며 성능이 붕괴됨을 확인했다.

UDM-GRPO는 GRPO의 그룹 상대 이득(Group Relative Advantage) 개념을 UDM에 이식했다. 한 프롬프트당 G개의 궤적을 생성하고 보상의 평균과 표준편차를 이용해 어드밴티지를 계산함으로써 별도의 가치 모델(Value Model) 없이도 효율적인 학습이 가능하다.

수학적으로는 정책 비율 r_t(θ)를 p_θ(x_1 | x_t) / p_old(x_1 | x_t)로 정의하여 최적화한다. Reduced-Step 전략은 전체 타임스텝 중 노이즈가 많은 전반부에서 연속된 3개의 스텝을 무작위로 선택해 학습함으로써 그라디언트 분산을 줄이고 수렴 속도를 높인다. 또한 CFG-Free 학습을 통해 조건부/비조건부 모델을 동시에 최적화해야 하는 복잡도를 제거하면서도 최종 성능에서 CFG 기반 방식보다 우수한 결과를 얻었다.

본 논문은 주로 텍스트-이미지 생성 및 OCR 작업에 집중하고 있으며, 비디오 생성이나 더 복잡한 멀티 리워드 최적화 작업에 대한 확장은 향후 과제로 남겨두고 있습니다.