Diffusion-State Policy Optimization for Masked Diffusion Language Models

Motivation

• Known: MDLM은 여러 denoising step에 걸쳐 마스크된 토큰을 반복적으로 채워 텍스트를 생성하며, diffu-GRPO나 SPG와 같은 terminal-feedback policy optimization 방법들이 tractable surrogate likelihood를 통해 group-based policy gradient(GRPO 계열)로 MDLM을 학습시키는 데 사용되어 왔다.
• Gap: 기존 terminal-feedback PO는 final completion에 대한 scalar reward 하나를 모든 중간 마스크 채우기 결정들에 동일하게 귀속시켜 credit assignment가 조악하며, 일부 선행연구가 reward shaping이나 고정 step에서의 trajectory-level 비교로 이를 densify하려 했으나 여전히 실현된 denoising trajectory 단위에 머물러 있어 상태-조건부(state-conditioned) 필링 행동 자체를 최적화 단위로 다루지 못했다.
• Why: MDLM은 매 denoising step마다 다수의 마스크 위치에 대한 token-level 분포를 anytime으로 제공하므로, 이 중간 정보를 추가 rollout이나 optimizer step 없이 재활용해 세밀한 credit assignment를 가능하게 하면 policy-gradient 기반 MDLM 학습의 표본 효율성과 성능을 동시에 개선할 수 있어 diffusion 기반 LLM의 RL 학습 일반에 실용적 가치가 크다.
• Approach: 중간 masked sequence를 state로, 그 시점의 mask filling을 action으로 정의하고, 선택된 중간 상태에서 rollout 시 캐시된 logits로부터 현재 마스크 위치들을 재샘플링(branching)하여 얻은 completion들을 동일한 terminal reward로 점수화한 뒤, 새로 채워진 토큰들에 대해서만 policy-gradient 업데이트를 적용하는 fixed-state objective를 terminal-feedback PO와 결합한다.

Achievement

Figure 4. Wall-clock-matched training curves on LLaDA-8B-

DiSPO는 LLaDA-8B-Instruct 모델에서 diffu-GRPO와 SPG라는 두 terminal-feedback MDLM PO 방법에 plug-in으로 결합되었을 때, 동일한 rollout 연산량과 optimizer step 수 조건 하에서 GSM8K, MATH500 등의 수학 추론 벤치마크와 Sudoku, Countdown 등의 symbolic planning 벤치마크에서 baseline 대비 일관된 성능 향상을 보였다.

How

Figure 3. Variance reduction of the step-wise gradients on

• 중간 masked state $s_{k,t}=(q,x_{k,t})$를 정의하고, 현재 마스크된 위치들의 결합 필링을 action $a_{k,t}$로 정의하며 factorized policy $\pi_\theta(a_{k,t}|s_{k,t})=\prod_{i\in M_{k,t}}\pi_\theta(a_{k,t,i}|s_{k,t},i)$로 표현
• 선택된 상태에서 rollout 시 캐시된 logits로부터 Z개의 alternative filling(drafts)을 샘플링해 branched completion을 생성하고 동일한 reward function으로 점수화
• fixed-state expected-return objective를 formalize하고 branched completion에 대한 policy-gradient estimator를 유도(Theorem A.1), terminal-feedback objective와 결합해 joint objective가 기댓값 상에서 성립함을 증명(Theorem A.2)
• 새로 채워진 토큰만 업데이트하고 same-state branch 간 평균을 취하는 것이 분산을 줄이고 step-wise update 효율을 높임을 이론적으로 뒷받침(Propositions A.3, A.4)
• diffu-GRPO 및 SPG를 base optimizer로 삼아 DiSPO를 Algorithm 1과 같이 joint loss $\alpha_{term}L_{term}+\alpha_{step}L_{step}$ 형태로 결합하여 실험

Originality

• 기존 MDLM policy optimization이 terminal reward만을 사용하는 bandit-style 관점에 머무른 것과 달리, 중간 masked state를 명시적 state로, mask filling을 action으로 정식화하여 policy gradient theorem의 세밀한 credit assignment 관점을 MDLM에 적용
• rollout 시 이미 계산된 logits를 재사용하는 same-state branching을 통해 추가 multi-step diffusion rollout이나 optimizer step 없이 credit assignment를 densify하는 계산 효율적 설계
• terminal-feedback PO(예: diffu-GRPO, SPG)를 대체 가능한 base optimizer로 취급하는 plug-in 형태로 설계하여 범용성 확보
• fixed-state objective에 대한 policy-gradient estimator 유도 및 joint objective의 기댓값 등가성, 그리고 새로 채워진 토큰만 업데이트/same-state branch 평균이 분산 감소에 기여함을 증명하는 이론적 뒷받침

Limitation & Further Study

• 실험이 LLaDA-8B-Instruct라는 단일 MDLM 백본과 diffu-GRPO, SPG 두 base optimizer 조합에 한정되어 있어 다른 MDLM 아키텍처나 더 큰 규모 모델로의 일반화 검증이 필요
• branch size Z와 중간 state 선택 전략(어떤 timestep에서 branching할지)이 성능에 미치는 민감도 분석이 본문 발췌만으로는 충분히 드러나지 않으며 이에 대한 추가적 ablation과 이론적 가이드가 필요
• surrogate likelihood $\tilde\pi_\theta$에 의존하는 근사적 특성이 DiSPO의 이론적 보장(fixed-state objective의 정확성)에 미치는 영향에 대한 심층 분석 필요
• 수학 추론 및 symbolic planning 외 open-ended generation이나 alignment 태스크로의 확장 가능성 및 reward hacking 위험에 대한 추가 검증이 후속 연구로 요구됨

Evaluation

총평: MDLM policy optimization에서 rollout 중 자연스럽게 생성되는 중간 상태 정보를 추가 비용 없이 재활용해 credit assignment를 개선한다는 아이디어가 이론적으로 잘 뒷받침되고 실험적으로도 일관된 개선을 보여주는 실용적이고 견고한 plug-in 기법이다.