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

2026-07-05

AI Summary

Diffusion language models struggle to enhance reasoning capabilities during post-training due to exposure bias in supervised fine-tuning and sparse rewards in reinforcement learning. This work proposes an online self-distillation strategy that leverages the model’s own denoising trajectories, using later, higher-quality delayed decoding states as teacher signals to provide dense, token-wise on-policy supervision—without requiring external labels. By extracting privileged information directly from the student model’s decoding process, the approach ensures alignment between distillation and inference behavior. Evaluated on the Dream and LLaDA benchmarks, the method significantly outperforms standard supervised fine-tuning and existing online strategies, achieving state-of-the-art performance in mathematical reasoning and code generation tasks.

Abstract

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.

Problem

• diffusion language models
• reasoning
• on-policy learning
• self-distillation
• post-training

Innovation

• diffusion language models
• on-policy self-distillation
• privileged information
• denoising trajectory
• token-level supervision