dOPSD: On-Policy Self-Distillation for Diffusion Language Models

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.

한 줄 요약

dOPSD는 확산 언어 모델의 자체 디스틸레이션을 통해 수학 추론과 코드 생성 성능을 향상시키는 on-policy 학습 방법이다.

핵심 기여도

기존 OPSD는 외부 레퍼런스 솔루션에 의존하여 확산 모델에서 효과적이지 못함을 밝힘.
dOPSD는 학습 중 생성된 디노이징 트레젝토리에서 선생 모델의 특권 정보를 유도함.
Dream-7B-Instruct와 LLaDA-8B-Instruct 데이터셋에서 수학 추론과 코드 생성 모두에서 기존 방법을 초과함.
롤아웃 검증 없이도 성능 개선이 가능함을 실험적으로 입증함.

핵심 아이디어

기존의 on-policy self-distillation(자체 디스틸레이션)은 외부 레퍼런스 솔루션을 기반으로 선생 모델의 특권 정보(PI)를 제공하지만, 이는 추론 시점에는 제공되지 않아 학습 효과가 제한된다. dOPSD는 이 문제를 해결하기 위해, 학습 중 생성된 디노이징 트레젝토리 내에서 선생 모델의 정보를 직접 추출한다. 즉, 학습 중에 생성된 더 많은 정보가 포함된 후속 단계의 컨텍스트를 활용하여, 마스킹된 위치에 대한 예측을 개선한다. 이는 확산 모델의 병렬 디코딩 특성에 기반한 새로운 통찰이며, 외부 레퍼런스 없이도 on-policy 학습이 가능하도록 한다.

기술적 접근법

**dOPSD**: 학습 중 생성된 디노이징 트레젝토리에서 선생 모델의 정보를 추출.
**Jensen–Shannon 목적 함수**: 선생 모델의 분포를 학생 모델에 디스틸.
**롤아웃 검증**: 최종 정답과 비교하여 신뢰할 수 있는 학습 신호만 사용.
**트레이닝 파이프라인**: 마스킹된 위치를 재마스킹하여 기존 확산 모델의 forward pass에 바로 적용 가능.
**하이퍼파라미터**: 특별히 언급되지 않음.
**데이터셋**: Dream, LLaDA.

주요 결과

**Dream-7B-Instruct**: 수학 추론 성능 개선.
**LLaDA-8B-Instruct**: out-of-domain 코드 생성 성능 개선.
기존 supervised, RL, naive OPSD 기반 방법 대비 **상승**.
롤아웃 검증 없이도 성능 개선 가능.
외부 레퍼런스 없이도 학습 가능.

의의 및 한계

dOPSD는 확산 언어 모델의 post-training에 있어 외부 레퍼런스에 의존하지 않는 새로운 학습 신호를 제시하며, 수학 추론과 코드 생성 모두에서 실질적인 성능 향상을 보인다. 특히, 기존 OPSD가 확산 모델에 적용되지 않는 이유를 분석하고 이를 해결한 점이 학술적으로 의미가 크다. 그러나, 본 연구는 수학 및 코드 생성에만 적용되었으며, 다른 도메인에서의 일반화 가능성은 추가 연구가 필요하다. 또한, 롤아웃 검증이 없어도 성능이 유지되지만, 이는 데이터셋의 특성에 따라 달라질 수 있다.

실용적 활용

dOPSD는 확산 언어 모델을 기반으로 하는 수학 문제 해결, 코드 생성, 병렬 텍스트 생성 등에 적용 가능하다. 특히, 레이블이 부족하거나 외부 레퍼런스를 사용할 수 없는 상황에서 유용하며, 산업적으로는 저비용으로 모델을 학습시키는 데 기여할 수 있다.