[논문리뷰] UDM-GRPO: Stable and Efficient Group Relative Policy Optimization for Uniform Discrete Diffusion Models

링크: 논문 PDF로 바로 열기

메타데이터

저자: Jiaqi Wang, Haoge Deng, Ting Pan, Yang Liu, Chengyuan Wang, Fan Zhang, Yonggang Qi, Xinlong Wang et al.

1. Key Terms & Definitions (핵심 용어 및 정의)

• UDM (Uniform Discrete Diffusion Model): 모든 토큰을 병렬로 업데이트하며 시간 의존적인 범주형 오염 과정을 명시적으로 매개변수화하는 이산 생성 모델입니다.
• GRPO (Group Relative Policy Optimization): 대규모 언어 모델의 추론 능력 강화를 위해 제안된 강화학습 기법으로, 그룹 내 보상을 정규화하여 정책을 최적화합니다.
• Trajectory (Forward/Backward): 확산 모델의 학습 및 추론 과정에서 정의되는 상태 경로로, 본 논문은 전방 확산 과정을 통한 X_forward 구성을 제안합니다.
• Action: 정책 최적화의 대상이 되는 값으로, 본 논문에서는 중간 예측치가 아닌 최종 정제된 샘플(Clean Sample, x_hat_1)을 액션으로 정의합니다.

2. Motivation & Problem Statement (연구 배경 및 문제 정의)

본 논문은 Uniform Discrete Diffusion 모델에 Reinforcement Learning을 안정적으로 통합하는 문제를 해결합니다. 기존의 Flow-GRPO를 UDM에 직접 적용할 경우, 훈련 과정에서 심각한 불안정성과 성능 저하가 발생한다는 한계가 있습니다. 이러한 불안정성은 초기의 부정확한 중간 예측치를 액션으로 사용하는 것과, 역방향 프로세스 기반의 학습이 사전 훈련 분포와 불일치(Distribution Shift)를 유발하는 데서 기인합니다 [Figure 1]. 따라서 모델의 훈련 안정성을 보장하고 사전 훈련 분포와의 정합성을 유지하는 새로운 통합 프레임워크가 필수적입니다.

3. Method & Key Results (제안 방법론 및 핵심 결과)

본 논문은 UDM과 GRPO를 안정적으로 통합하기 위해 UDM-GRPO 프레임워크를 제안합니다. 첫째, 모든 타임스텝에서 액션을 중간 예측치가 아닌 최종 정제 샘플 x_hat_1으로 재정의하여 보상 일관성과 최적화 정밀도를 높였습니다 [Figure 3]. 둘째, 역방향 프로세스 대신 전방 확산 과정을 사용하여 학습 궤적(X_forward)을 재구성함으로써, 사전 훈련 분포와의 정합성을 강화하고 OOD 문제를 해결하였습니다. 또한, 효율성을 위해 고노이즈 타임스텝에 집중하는 Reduced-Step 전략과 CFG-Free 학습 방식을 도입하였습니다 [Figure 3]. 정량적 실험 결과, GenEval 정확도는 69%에서 96%로 비약적으로 향상되었으며, PickScore 역시 20.46에서 23.81로 개선되어 연속 및 이산 설정 모두에서 SOTA 성능을 달성했습니다 [Table 2]. 특히, OCR 벤치마크에서는 8%에서 57%로 정확도가 상승하여 모델의 강건한 일반화 능력을 입증했습니다 [Table 2].

4. Conclusion & Impact (결론 및 시사점)

본 연구는 Uniform Discrete Diffusion 모델에 강화학습을 적용하는 최초의 프레임워크인 UDM-GRPO를 성공적으로 구축하였습니다. 액션 재정의 및 전방 궤적 재구성을 통해 기존 RL 통합 방식의 근본적인 불안정성 문제를 효과적으로 해결하였습니다. 이러한 성과는 향후 이산 확산 모델의 인간 피드백 기반 정렬(Alignment) 분야에서 새로운 표준을 제시하며, 텍스트-투-이미지 및 멀티모달 생성 모델의 품질 향상에 크게 기여할 것으로 기대됩니다.

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.