dOPSD: 확산 언어 모델을 위한 온-정책 자기 증류

초록

확산 대규모 언어 모델(dLLM)은 마스킹된 시퀀스를 반복적으로 잡음 제거하여 텍스트를 생성하며, 자기회귀 모델에 대한 병렬적 대안을 제공한다. 그러나 사후 훈련을 통해 강력한 추론 능력을 유도하는 것은 여전히 어렵다. 지도 미세 조정은 오프-정책(off-policy)이며 노출 편향(exposure bias)을 겪는 반면, 강화 학습은 드문 시퀀스 수준의 보상만을 제공하고 계산 가능한 시퀀스 우도 없이는 적용하기 어렵다. 온-정책 자기 증류(OPSD)는 하나의 모델을 학생이자 교사로 사용하여 밀집된 토큰 수준의 온-정책 감독을 제공하는 유망한 대안을 제시하지만, 그 효과는 교사에게 특권 정보(PI)를 부여하는 데 달려 있다. 일반적으로 이는 추론 시 사용할 수 없는 인스턴스별 실제 정답 참조인데, 그 결과 학생은 PI 없는 약한 합의 정책을 증류하게 되어 dLLM 추론에 거의 개선을 가져오지 못한다. 우리는 dOPSD를 도입하며, 이는 교사의 특권을 학생 자신의 잡음 제거 궤적으로부터 직접 도출한다. 즉, 외부 레이블 대신 동일한 궤적의 이후의 더 많이 디코딩된 단계를 사용하여 마스킹된 위치를 평가함으로써, 교사의 이점이 모델 자체의 디코딩 과정에서 발생하게 한다. Dream과 LLaDA에서 dOPSD는 도메인 내 수학적 추론과 도메인 외 코드 생성을 모두 개선하며, 지도 및 온-정책 기준선을 능가한다.

English

Diffusion large language models (dLLMs) generate text by iteratively denoising a masked sequence, offering a parallel alternative to autoregressive models, but eliciting strong reasoning through post-training remains difficult: supervised fine-tuning is off-policy and suffers from exposure bias, while reinforcement learning gives only sparse, sequence-level rewards and is hard to apply without tractable sequence likelihoods. On-policy self-distillation (OPSD) offers a promising alternative, using one model as both student and teacher to provide dense, token-level, on-policy supervision, but its effectiveness hinges on giving the teacher privileged information (PI) - typically an instance-specific ground-truth reference unavailable at inference - so the student ends up distilling a weak PI-free consensus policy that yields little improvement on dLLM reasoning. We introduce dOPSD, which instead derives the teacher's privilege directly from the student's own denoising trajectory, evaluating masked positions using later, more-decoded steps of that same trajectory rather than an external label, so the teacher's advantage emerges from the model's own decoding process; on Dream and LLaDA, dOPSD improves both in-domain math reasoning and out-of-domain code generation, outperforming supervised and on-policy baselines.