Diffusion language models can learn from autoregressive teachers: JustOPD lifts Dream-7B's math reasoning while keeping any-order decoding
Synopsis
Autoregressive models generate strictly left to right, while diffusion language models denoise in arbitrary order, so the former's token-level supervision does not transfer directly. JustOPD forces the diffusion student to roll out left to right during training, letting a Qwen3-14B teacher supervise every intermediate state, while inference keeps any-order decoding. With Dream-7B as student and Qwen3-14B as teacher on GSM8K and MATH500, scores rise over the base model across several divergences.
Interpretation
A diffusion language model can be on-policy distilled directly from a stronger autoregressive language model, without first finding an equally strong diffusion teacher. Prior on-policy distillation for diffusion models was mainly diffusion-to-diffusion, and strong open diffusion teachers are scarce; autoregressive teachers are plentiful, but their left-to-right decoding does not match the diffusion student's arbitrary-order denoising. With Dream-7B as student and Qwen3-14B as teacher on GSM8K and MATH500, JustOPD beats the base Dream-7B under every tested divergence, and evaluation uses the diffusion model's native any-order decoder.
Making the diffusion student generate token by token from left to right during training turns every intermediate state the teacher sees into a valid prefix. Mid-denoising, a diffusion model holds a partially masked sequence that generally does not correspond to a valid prefix under the autoregressive teacher's training distribution, so the teacher cannot supply meaningful token-level supervision. The student policy is written as predicting the clean token at the next position given a prefix plus a masked suffix, matching the autoregressive teacher's trajectory factorization so the standard on-policy distillation objective applies; the implementation uses a sampled-token, PPO-style objective and computes loss only at the teacher's highest-entropy positions.
The training-time ordering constraint does not change how the model generates at inference, where the student still decodes in any order. Carrying left-to-right generation into inference would give up the parallel generation and lower latency that diffusion models offer. Evaluation runs Dream-7B with an entropy-based unmasking strategy that unmasks the two lowest-entropy positions per step, the diffusion model's native any-order decoding; knowledge from distillation still shows up as higher scores under that decoder.
Perspective
Experiments use Dream-7B as the student and Qwen3-14B, plus task-specific fine-tuned versions of it, as teachers, evaluated on the GSM8K and MATH500 math reasoning benchmarks across several divergences and generation lengths; distillation runs for 100 training steps, with evaluation every 10 steps and the best checkpoint reported.
The authors note that current token-level distillation assumes teacher and student share a tokenizer, limiting which model pairs can be combined; AIME 2024 results are preliminary, and the base model solves no problems there even at longer generation lengths; cross-tokenizer distillation and extension to block-attention diffusion models remain future work.
