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

The Core Problem

Diffusion language models—which generate text by gradually unmasking random noise rather than predicting tokens left-to-right—are theoretically attractive (they can run in parallel) but notoriously hard to improve after training. Standard approaches all fail in different ways. Supervised fine-tuning ignores the model's own behavior during inference, creating a train-test mismatch. Reinforcement learning requires tractable probability scores for entire sequences, which diffusion models don't naturally provide. The authors identify a fundamental weakness in existing on-policy self-distillation for diffusion models: the teacher typically needs access to a ground-truth answer or external signal to outperform the student, but this privilege vanishes at test time, so the student learns little.

The Key Insight

dOPSD solves this by making the teacher's advantage self-contained. Instead of relying on an external reference answer, the teacher evaluates masked token positions by peeking at later steps of the model's own denoising trajectory. In other words: if the model has already decoded positions 5, 10, and 15 in the current pass, the teacher uses that partial context to score positions 3 and 7, which the student hasn't seen yet. This creates genuine privileged information—the teacher knows more of the final output—without requiring a human label or gold answer. Critically, this information gap can be closed during inference by simply running the model longer, so the student's learned policy directly transfers.

Why It Works and Matters

The method produces two concrete improvements. On mathematical reasoning tasks (Dream benchmark), dOPSD outperforms both supervised baselines and simpler on-policy methods. More tellingly, it also generalizes to code generation (LLaDA benchmark), demonstrating that the learned reasoning capability isn't overfit to the training domain.

The contribution is subtle but important. Diffusion models' parallel decoding is their main advantage over autoregressive models, yet training them to reason has been stuck between weak supervised approaches and intractable RL setups. dOPSD bridges this gap by extracting training signal directly from the model's own iterative refinement process—no external data, no sparse rewards, no likelihood computations. The teacher's edge is transient (it only exists during training), but sufficient to guide the student toward better policies that remain effective at inference.

For practitioners building reasoning systems with diffusion LLMs, this offers a practical path to post-training improvement that respects the model's architecture rather than fighting it.