확산 언어 모델을 위한 효율적 정책 그라디언트: AGRPO

초록

본 논문은 확산 기반 대형 언어 모델(dLLM)의 사후 학습에 적용 가능한 새로운 정책 그라디언트 알고리즘인 Amortized Group Relative Policy Optimization(AGRPO)를 제안한다. 다단계 마코프 의사결정 과정을 이용해 개별 디노이징 스텝을 최적화함으로써, 기존의 ELBO 기반 근사법이 갖는 편향을 제거하고, 수학·논리 추론 벤치마크에서 크게 향상된 성능과 4배 빠른 샘플링을 달성한다.

상세 분석

이 논문은 확산 언어 모델(dLLM)이 전통적인 자동회귀(AR) LLM과 달리 시퀀스 수준의 정확한 확률을 계산할 수 없다는 근본적인 한계를 정확히 짚어낸다. 기존 연구들은 ELBO 혹은 기타 하한을 사용해 근사적인 로그우도 값을 정책 그라디언트에 대입했지만, 이는 편향된 업데이트를 초래한다. 저자들은 dLLM의 생성 과정을 “부분 마스크 → 토큰 복원”이라는 다단계 마코프 의사결정 프로세스로 모델링하고, 각 디노이징 스텝을 하나의 행동으로 보는 새로운 MDP 정의를 제시한다. 이 접근법은 각 스텝에서의 행동 확률을 정확히 계산할 수 있게 해, REINFORCE 혹은 PPO와 같은 기존 정책 그라디언트 공식에 그대로 적용할 수 있다.

핵심 기여는 두 가지이다. 첫째, “시간 샘플링” 기법을 도입해 전체 m 스텝을 모두 전방패스하지 않고도 무편향 추정량을 얻는다. 구체적으로, 스텝 t를 균등하게 샘플링하고, 해당 스텝의 중요도 비율 ρₜ와 그룹 정규화된 어드밴티지를 사용해 기대값을 근사한다. 이는 메모리와 연산 비용을 크게 절감하면서도 편향을 없앤다. 둘째, KL 발산을 포함한 클리핑 손실을 그대로 유지하면서, 샘플링된 스텝에 대해 Monte‑Carlo 추정량을 사용해 분산을 감소시키는 두 가지 변형(베이스라인 보정, 스트라티파이드 샘플링)을 제안한다. 이들 기법은 대규모 트랜스포머 모델에서도 안정적인 학습을 가능하게 한다.

실험에서는 GSM8K, MATH‑500, Countdown, Sudoku 네 가지 추론·수학 벤치마크에 AGRPO를 적용하였다. 기존 dLLM 사후 학습 방법인 diffu‑GRPO 대비 각각 +9.9%, +4.6%, +59.4%, +69.7%의 절대 정확도 향상을 기록했으며, 특히 Sudoku와 Countdown 같은 복합 논리 문제에서 현저한 개선을 보였다. 또한, 사후 학습된 모델은 샘플링 스텝 수를 4배까지 줄여도 성능 저하가 미미함을 확인했는데, 이는 AGRPO가 학습 과정에서 다양한 스텝 수에 대한 일반화를 내재화했기 때문이다.

이 논문은 dLLM이 기존 AR LLM과 동등하거나 그 이상으로 추론 능력을 발휘하도록 만드는 중요한 전환점을 제공한다. 다단계 MDP와 무편향 정책 그라디언트라는 이론적 토대 위에 실용적인 메모리 절감 및 분산 감소 기법을 결합함으로써, 확산 기반 언어 모델이 RL‑VR(Verification‑Reward) 설정에서도 강력히 활용될 수 있음을 입증한다. 향후 연구는 더 복잡한 보상 구조(예: 단계별 피드백)와 다른 디노이징 스케줄에 대한 확장, 그리고 대규모 멀티모달 dLLM에 대한 적용 가능성을 탐색할 여지를 남긴다.