[Paper] UDM-GRPO: Uniform Discrete Diffusion Models를 위한 안정적이고 효율적인 Group Relative Policy Optimization

Source: arXiv - 2604.18518v1

개요

Uniform Discrete Diffusion Models (UDMs)는 텍스트, 토큰, 이미지 패치와 같은 이산 데이터를 생성하는 데 있어 대표적인 도구가 되었습니다. 확산 기반 생성이 연속 영역에서 크게 발전한 반면, 이를 강화 학습(RL)과 결합하는 것은 아직 해결되지 않은 문제였습니다. 본 논문에서는 UDM‑GRPO를 소개합니다. 이는 UDM과 Group Relative Policy Optimization (GRPO)을 결합한 최초의 프레임워크로, 안정적인 학습 파이프라인을 제공하고 텍스트‑투‑이미지(T2I) 및 OCR 작업에서 큰 성능 향상을 달성합니다.

주요 기여

• 첫 번째 RL‑기반 UDM 프레임워크 – 이산 확산과 정책‑그라디언트 방법 사이의 격차를 메웁니다.
• 행동 중심 형식화 – 최종 깨끗한 샘플을 RL “행동”으로 간주하여 더 명확하고 분산이 낮은 그라디언트를 제공합니다.
• 전방 확산 과정을 통한 궤적 재구성 – 학습 신호를 사전 훈련 중에 본 분포와 정렬합니다.
• 두 가지 효율성 트릭 – Reduced‑Step (RL 미세조정 중 확산 단계 수를 줄임) 및 CFG‑Free (분류기‑프리 가이던스 오버헤드 제거) 로 품질을 손상시키지 않으면서 훈련 시간을 단축합니다.
• 최첨단 결과 – GenEval 정확도를 69 %에서 96 %로, PickScore를 20.46에서 23.81로 향상시키고; OCR 정확도는 8 %에서 57 %로 급증합니다.

Methodology

1. Baseline UDM – 사전 학습된 이산 확산 모델로, 잡음이 섞인 토큰 시퀀스를 반복적으로 디노이징하여 깨끗한 출력이 나올 때까지 진행합니다.
2. RL 문제 정의
   • State: 특정 확산 단계에서의 중간 잡음 토큰 시퀀스.
   • Action: 전체 역확산이 끝난 후 생성되는 최종 깨끗한 토큰 시퀀스.
   • Reward: 작업별 메트릭(예: T2I의 경우 CLIP 기반 유사도, OCR 정확도 등).
3. Group Relative Policy Optimization (GRPO) – 현재 정책을 과거 정책들의 그룹과 비교하여 분산을 줄이고 업데이트를 안정화시키는 정책 그래디언트 알고리즘.
4. Forward‑process trajectory reconstruction – 역궤적을 직접 샘플링하는 대신, 저자는 깨끗한 샘플에서 잡음이 섞인 상태로 전방 확산을 수행합니다. 이를 통해 샘플링된 궤적이 사전 학습된 모델이 학습된 동일한 매니폴드 위에 있음을 보장합니다.
5. Efficiency enhancements
   • Reduced‑Step: RL 파인튜닝 중에 확산 체인을 축소(예: 100 → 20 단계)하여 계산량을 크게 낮춥니다.
   • CFG‑Free: classifier‑free guidance가 필요 없게 하여 손실을 단순화하고 추론 속도를 높입니다.

전체 학습 루프는 후보 출력을 생성하고, 그 보상을 평가한 뒤, 재구성된 전방 궤적을 고려한 GRPO 그래디언트로 확산 정책을 업데이트하는 과정을 반복합니다.

결과 및 발견

• 안정성 – 학습 곡선이 부드럽게 수렴하는 반면, UDM에 대한 순수 GRPO는 몇 에포크 후에 발산합니다.
• 샘플 품질 – 정성적 예시에서 더 선명하고 의미적으로 정렬된 이미지와 더 깔끔한 토큰 시퀀스를 확인할 수 있습니다.
• 효율성 – Reduced‑Step은 최종 점수 손실이 거의 없으면서 RL 미세조정 시간을 약 70 % 단축합니다.

실용적 함의

• Diffusion 모델의 더 나은 파인튜닝 – 개발자들은 이제 RL을 적용해 작업별 목표(예: 생성된 이미지를 목표 스타일에 맞추기)를 개선할 수 있으며, 처음부터 재학습할 필요가 없습니다.
• 낮은 연산 비용 – Reduced‑Step 트릭 덕분에 RL 기반 정제가 단일 GPU에서도 가능해져, 스타트업 및 연구실에서 빠른 프로토타이핑이 가능해집니다.
• 크로스모달 생성 – 이 프레임워크는 이미지와 텍스트 생성 모두에 적용 가능하며, 코드 합성, 음성 토큰 생성, 게임 레벨 디자인 등 다른 이산 도메인으로의 손쉬운 확장을 시사합니다.
• 플러그‑앤‑플레이 – 이 방법이 사전 학습된 UDM을 기반으로 하기 때문에, 기존 파이프라인(예: 이산화된 Stable Diffusion 유사 모델)을 몇 줄의 RL 코드만으로 업그레이드할 수 있습니다.

제한 사항 및 향후 연구

• 보상 설계 의존성 – 성능은 신뢰할 수 있고 미분 가능한 보상(예: CLIP 점수)에 크게 좌우됩니다. 부실한 보상은 여전히 최적이 아닌 정책을 초래할 수 있습니다.
• 감소‑스텝 트레이드‑오프 – 학습 속도는 빨라지지만, 지나치게 공격적인 스텝 감소는 모델이 장거리 의존성을 탐색하는 능력을 제한할 수 있습니다.
• 대규모 어휘에 대한 확장성 – 논문에서는 중간 규모 토큰 집합을 대상으로 평가했으며, 매우 큰 어휘(예: 전체 문장 수준 언어 모델)로 확장하려면 추가적인 메모리 최적화 기법이 필요할 수 있습니다.
• 향후 방향 – 저자들은 계층적 확산 스텝, 자동 보상 형태화 탐색, 그리고 이산 행동 공간을 갖는 로봇공학 등 멀티모달 강화 학습 시나리오에 UDM‑GRPO를 적용하는 연구를 제안합니다.

저자

• Jiaqi Wang
• Haoge Deng
• Ting Pan
• Yang Liu
• Chengyuan Wang
• Fan Zhang
• Yonggang Qi
• Xinlong Wang

논문 정보

• arXiv ID: 2604.18518v1
• 분류: cs.CV, cs.LG
• 발표일: 2026년 4월 20일
• PDF: Download PDF