PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Phuong Tuan Dat, Qi Li, Xinchao Wang

arXiv:2607.04428v1cs.CLcs.AI

TL;DR

Strong reasoning post-training for diffusion language models lacks dense, on-policy supervision without external privileged information. dOPSD derives that privilege from later, more-decoded states in the student’s own denoising trajectory, achieving the best results across in-domain math and out-of-distribution code benchmarks on Dream and LLaDA.

  • Problem

    Diffusion language models need dense, on-policy, external-teacher-free post-training signals because sparse rewards and unavailable sequence likelihoods hinder reinforcement learning, while OPSD relies on unavailable reference-based privilege.

  • Method

    dOPSD uses later, more-decoded states from the student’s own denoising trajectory as privileged teacher contexts for supervising earlier masked positions.

  • Results

    dOPSD achieves the best score on every in-domain and out-of-distribution benchmark across Dream-7B-Instruct and LLaDA-8B-Instruct, raising Dream GSM8K from 81.41 to 83.04.

  • Takeaways & Limitations

    The results support deriving privileged information from diffusion decoding itself as an effective alternative to external reference solutions for dLLM post-training.

Abstract

from arXiv · show

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.

1 Introduction

dOPSD adapts on-policy self-distillation to diffusion language models by deriving teacher privilege from later states in the student’s own denoising trajectory. This avoids external references while providing dense token-level supervision for masked positions, addressing the limitations of sparse-reward reinforcement learning and off-policy supervised fine-tuning.

  • Motivation: Diffusion language models generate text through iterative denoising from a fully masked sequence, providing progressively more informed intermediate contexts.At each step, the model predicts masked tokens and commits the most confident ones.
  • Motivation: Sparse sequence-level rewards and intractable sequence likelihoods make reinforcement learning difficult to apply to diffusion language models.This motivates a dense, on-policy signal that does not require external teachers or reward models.
  • Prior approach: OPSD uses a student conditioned only on the problem and a teacher additionally conditioned on a ground-truth reference to provide dense per-token supervision along the student’s rollout.The method was proposed for autoregressive LLMs and depends on privileged information unavailable at inference.
  • Core idea: Later denoising states provide a peek-ahead advantage because they reveal tokens and context that earlier states must predict while still masked.The teacher evaluates positions from a more-unmasked trajectory state, making its privilege intrinsic to decoding.
  • Method: dOPSD derives privileged information from the student’s denoising trajectory and distills the later-state teacher distribution into earlier masked positions with a token-level Jensen–Shannon objective.The student records its completion history, and intermediate states are formed by re-masking undecoded positions.
  • Contribution: The original OPSD recipe fails to yield reliable diffusion-model gains because its external reference solution is unavailable at inference and produces a weak PI-marginalized consensus.dOPSD retains the privileged-teacher principle while sourcing privilege from the student’s own trajectory.

2 Related Work

Related work establishes diffusion language models as bidirectional, confidence-ordered alternatives to left-to-right generation and positions dOPSD among distillation methods that address off-policy mismatch. Prior OPSD uses ground-truth privileged information for self-teaching, whereas on-policy distillation generally relies on a separate larger teacher.

  • Diffusion language models: Masked diffusion models reverse progressive token corruption by predicting clean tokens at masked positions, enabling bidirectional context and parallel, confidence-ordered decoding.LLaDA and Dream scale this paradigm to billions of parameters.
  • Distillation: Knowledge distillation transfers a teacher’s soft predictions to a student, but fixed-corpus training is off-policy and suffers train-inference mismatch.The student is supervised on prefixes it would not itself produce.
  • On-policy distillation: On-policy distillation trains on student-generated sequences, providing dense feedback where the student actually goes and addressing compounding errors.Existing approaches assume a separate, larger teacher, unlike the same-model setup used here.
  • Privileged-information self-distillation: OPSD makes a model its own teacher by conditioning it on ground-truth privileged information and distilling per-token distributions into the student during on-policy rollouts.This builds on the intuition that explaining a known answer is easier than producing one.

3 Method

dOPSD adapts on-policy self-distillation to diffusion language models by using genuine denoising states and later states from the same student trajectory as teacher privilege. It addresses PI-free collapse and off-path random masking through trajectory-derived token-level supervision, with optional verification of correct rollouts.

  • OPSD: OPSD uses one model as an inference-conditioned student and privileged teacher, producing dense on-policy token-level supervision through a divergence objective.The teacher reads privileged information unavailable to the student, while gradients flow only through the student.
  • Diffusion decoding: For dLLMs, decoding is a confidence-ordered denoising trajectory whose masked positions shrink monotonically as tokens are committed.At each intermediate state, the model predicts tokens only at positions that remain masked.
  • Failure modes: Applying OPSD directly causes PI-free collapse on instance-specific references and trains on uniformly remasked states that do not occur along the model’s decoding path.These obstacles weaken reasoning because random masking ignores the confidence-ordered, easy-to-hard inference schedule and its committed context.
  • dOPSD: dOPSD samples a substantially masked intermediate trajectory state for the student and averages later predictions of the same still-masked positions as the teacher target.Later states provide more committed surrounding context, while positions already decoded are excluded from teacher averaging.
  • dOPSD: dOPSD applies token-level generalized Jensen–Shannon distillation only to currently masked positions, adds no architectural change, and can discard incorrect rollouts during training.Correctness verification uses only the short final answer; without reference answers, the method can distill from every rollout.

4 Experiments

Experiments on Dream-7B-Instruct and LLaDA-8B-Instruct show that dOPSD improves mathematical reasoning and transfers to out-of-distribution code generation, outperforming supervised, reinforcement-learning, and naive OPSD baselines. Ablations further favor forward KL, longer teacher horizons, a moderate mask threshold, and training without rollout verification.

  • Experimental setup: Experiments use Dream-7B-Instruct and LLaDA-8B-Instruct, training on MixChain-Z-PRM12K mathematical reasoning and evaluating transfer to code generation.Evaluation covers GSM8K and MATH500 accuracy, plus HumanEval and MBPP pass@1, without further code fine-tuning.
  • Main results: dOPSD is the only method improving the base model across both backbones, achieving the best score on every in-domain and out-of-distribution benchmark.On Dream, GSM8K rises from 81.41 to 83.04 and MATH500 from 38.97 to 42.20; on LLaDA, MATH500 improves by 4.76 points and GSM8K by 1.64.
  • Main results: SFT, GRPO, and naive OPSD baselines fail to improve reasoning reliably, while full-solution OPSD loses 14.2 GSM8K points on Dream and 13.4 on LLaDA.Answer-only OPSD remains near the base model, whereas GRPO’s only reported gain is +1.53 on LLaDA MATH500.
  • Out-of-distribution transfer: dOPSD transfers to code generation, lifting HumanEval by 4.17 points on Dream and 3.51 on LLaDA, while SFT and both OPSD ports lose ground.The reported changes are Dream 52.54 →56.71 and LLaDA 36.12 → 39.63.
  • Ablations: Forward KL outperforms reverse KL on every task-backbone cell, including LLaDA HumanEval rising from 37.42 to 39.63 and Dream MBPP from 50.61 to 58.49.The comparison uses otherwise identical settings, with forward KL as the default β →0 case and reverse KL as β →1.
  • Ablations: Performance improves with longer teacher horizons, and τ = 0.5 gives the highest four-task average at 60.11; without rollout verification, the average still rises from 57.59 to 58.64.Verification adds a further boost, while dOPSD can train from every rollout when reference answers are unavailable.

5 Conclusion

dOPSD adapts on-policy self-distillation to diffusion language models by addressing why the autoregressive OPSD recipe does not transfer. It replaces unavailable external privileged information and off-path random masking with mechanisms derived from diffusion-model decoding.

  • dOPSD adapts the privileged-teacher principle of OPSD to diffusion language models.
  • Autoregressive OPSD fails to transfer because its privileged information is an external, instance-specific reference solution unavailable at inference.This causes OPSD to collapse onto a weak PI-free consensus.
  • Random masking also fails because it places training states off the model’s confidence-ordered decoding path.

A Dataset Details

This appendix provides expanded details on the training corpus and evaluation benchmarks summarized in the Experimental Setup, covering their sources, sizes, preprocessing, and scoring protocols.

  • Scope: The appendix expands the Experimental Setup’s summary of the training corpus and evaluation benchmarks.It provides additional dataset-level information beyond the summary.
  • Data sources and scale: It documents the sources and sizes of the training corpus and evaluation benchmarks.
  • Processing and evaluation: It also describes preprocessing and scoring protocols for the corpus and benchmarks.

A.1 Training Data

All methods are post-trained on MixChain-Z-PRM12K, a roughly 12K-problem mathematical-reasoning corpus pairing checkable answers with chain-of-thought solutions of varied lengths. The final answer supports answer-level verification, while reference solutions are reserved for baselines and dOPSD avoids using them as privileged information.

  • A.1 Training Data: MixChain-Z-PRM12K contains roughly 12K competition-style mathematics problems paired with checkable final answers and chain-of-thought solutions of varying reasoning length.The corpus follows the MATH/PRM800K lineage and includes terse through fully elaborated solutions for the same problem.
  • A.1 Training Data: The rollout verifier uses only each problem’s gold final answer for answer-level checking and never exposes the reference chain of thought.SFT and full-solution OPSD use reference solutions, whereas dOPSD does not consume them as privileged information.
  • A.1 Training Data: All methods use the same instruction-style prompt template for training and evaluation, discard unparsable gold answers, and tokenize with each backbone’s native tokenizer.Completions are truncated to the model’s generation length, with a 256-token completion budget matching diffusion decoding length.

A.2 Evaluation Benchmarks

The evaluation spans four non-overlapping public benchmarks: GSM8K and MATH500 for in-domain mathematical reasoning, plus HumanEval and MBPP for out-of-distribution code generation. Math uses Acc@1 final-answer accuracy, while code uses pass@1 functional correctness, with single-sample greedy decoding.

  • Benchmark suite: The benchmark suite contains two mathematical-reasoning datasets and two code-generation datasets, with code testing transfer to a domain unseen during post-training.None of the benchmarks overlaps with the training corpus.
  • Mathematical reasoning: GSM8K contains 1,319 grade-school arithmetic problems requiring multi-step numerical reasoning and uses exact-match final-answer Acc@1.A completion is correct when its extracted final answer exactly matches the gold value.
  • Mathematical reasoning: MATH500 is a 500-problem competition-mathematics subset spanning seven subjects and five difficulty levels, scored by mathematically equivalent answers using Acc@1.Equivalent surface forms of the same value count as correct.
  • Code generation: HumanEval comprises 164 hand-written Python problems with hidden unit tests, and correctness requires generated functions to pass all tests under pass@1.Each problem provides a function signature and natural-language docstring.
  • Code generation: MBPP uses a sanitized 500-problem test split of entry-level Python tasks and reports pass@1 under the same functional-correctness criterion as HumanEval.Each task includes a prompt and a small suite of unit tests.

B Hyperparameters for Inference · C Qualitative Examples

Inference uses single-sample decoding with backbone- and benchmark-specific configurations, including one-token-per-step diffusion and tailored remasking or block settings. Qualitative examples expose distinct GSM8K baseline errors that dOPSD avoids.

  • B Hyperparameters for Inference: All evaluations use single-sample decoding, reporting Acc@1 for math and pass@1 for code without best-of-n or self-consistency.Table 5 provides the complete backbone- and benchmark-specific configurations.
  • B Hyperparameters for Inference: Diffusion steps equal generated tokens across backbones because committing multiple tokens per step sharply degrades quality.Generation lengths are benchmark-specific, covering reference solutions within each model’s position limit; code uses longer lengths than math.
  • B Hyperparameters for Inference: Dream uses entropy remasking at temperature 0.1 with top-p = 0.9 and no top-k truncation.HumanEval additionally uses entropy-threshold parallel decoding with threshold 0.5; other Dream runs use the plain entropy sampler.
  • B Hyperparameters for Inference: Table 5 specifies inference hyperparameters across base models and benchmarks, with diffusion steps equal to generated tokens on both backbones.A dash denotes an inapplicable or unused hyperparameter; for LLaDA, temperature is 0 and top-p is informational only.
  • B Hyperparameters for Inference: LLaDA uses lowconfidence remasking at temperature 0, ignores top-p = 0.9 under Gumbel-argmax sampling, and varies semiautoregressive block length by benchmark.Block length ranges from 8 on GSM8K to one full-length block on HumanEval.
  • C Qualitative Examples: Three representative Dream-7B-Instruct GSM8K generations illustrate distinct baseline errors that dOPSD avoids, with highlighting marking pivotal correct and erroneous steps.The examples make baseline failure modes concrete.

C.1 Example 1: Problem 526

In Problem 526, dOPSD is the only method that avoids an unnecessary ×24 conversion and reaches the gold answer of 21. All baselines make unit or subtraction errors, while the full-solution variant produces an unfinished malformed answer.

  • Model comparison: 21 is dOPSD’s answer, and it is the only method that reads the units correctly without inserting a spurious ×24 hours conversion.Every baseline hallucinates the conversion, with some also reversing the bill subtraction.
  • Correct solution: $21 is the correct difference: the original weekly bill is $126 and the new weekly bill is $147.The calculation uses daily rates directly and multiplies by 7 days.
  • Model comparison: dOPSD’s trajectory-derived teacher preserves unit-consistent reasoning, whereas supervised, reinforcement-learning, and PI-conditioned baselines inherit or amplify the ×24 error.The full-solution variant additionally rambles into an unfinished, malformed answer.

C.2 Example 2: Problem 589

In Problem 589, Martin loses 6.25 pounds on Cheerios and gains 8.75 pounds on donuts over five weeks, so the signed difference is 15 pounds. dOPSD gives the gold answer by preserving the opposite signs, whereas baselines return 2.5 by subtracting magnitudes, and one response is empty.

  • Correct solution: 15 pounds is the correct difference because the loss and gain lie on opposite sides of Martin’s starting weight.The signed calculation is 8.75−(−6.25) = 15.
  • Model comparison: dOPSD alone recognizes that the loss and gain add together and returns the gold answer 15.Every baseline instead computes 8.75−6.25 = 2.5 by silently dropping the sign.
  • Problem setup: Over 5 weeks, Cheerios causes a 6.25-pound loss while donuts causes an 8.75-pound gain.The weekly changes are −1.25 and +1.75 pounds over 5 weeks.
  • Model comparison: Several baseline responses report 2.5 pounds by subtracting the two magnitudes.This result appears as 8.75−6.25 = 2.5 in the baseline passages.
  • Model outputs: One model response is empty and generates no text.Its raw output is recorded as [’ ’].

C.3 Example 3: Problem 780

In Problem 780, Ava’s score change is −25 relative to Emma, so the difference in scores is the positive magnitude 25. dOPSD alone reports 25, while baselines retain or mishandle the signed gap and the full-solution teacher can produce malformed or empty outputs.

  • Method comparison: Only dOPSD reports the requested difference as the positive magnitude 25; baselines output −25, claim it cannot be found, or collapse to 0.The baselines either box the signed value, second-guess the unknown absolute scores, or misinterpret the requested magnitude.
  • Problem 780: Ava’s net score change is 80 + 15 − 120 = −25, making her score 25 points lower and the score difference 25.The calculation uses 8×10 = 80, 3×5 = 15, and 4×30 = 120.
  • Method comparison: The full-solution teacher degrades into malformed or empty generations, including an empty response.One listed output contains no generated text at all.
Loading 2607.04428v1…