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

Phuong Tuan Dat, Qi Li, Xinchao Wang

dOPSD improves diffusion language model reasoning by distilling the model's own future denoising steps as privileged information.

How can we adapt on-policy self-distillation (OPSD) to diffusion language models (dLLMs) to improve their reasoning capabilities?

Diffusion language models (dLLMs) struggle with complex reasoning because standard post-training methods like supervised fine-tuning suffer from exposure bias, while reinforcement learning provides only sparse, sequence-level feedback. dOPSD addresses this by using the model's own denoising trajectory as a teacher: at any intermediate step, the model acts as a student predicting masked tokens, while the same model at later, more-decoded steps acts as a privileged teacher that has already "peeked ahead" at the solution. This approach provides dense, on-policy, token-level supervision without requiring external labels or architectural changes, consistently outperforming supervised and reinforcement-learning baselines on mathematical and coding benchmarks.

Paper Primer

The core mechanism hinges on the iterative nature of diffusion decoding. As a dLLM fills in masked tokens, later states in the trajectory naturally contain more information about the final sequence than earlier ones; dOPSD treats these later states as a "privileged" teacher that guides the student's predictions at earlier, more-masked steps.

dOPSD consistently improves reasoning performance across both in-domain math and out-of-distribution code generation.

On Dream-7B-Instruct, dOPSD raised GSM8K accuracy from 81.41 to 83.04 and HumanEval pass@1 from 52.54 to 56.71. It is the only tested post-training method that improves the base model across all evaluated benchmarks.

The method avoids the "PI-free collapse" seen in previous on-policy self-distillation (OPSD) attempts. By sourcing privilege from the model's own trajectory rather than an external, instance-specific reference solution, dOPSD prevents the student from averaging over incompatible teachers and instead distills a coherent, self-generated signal.

Why does dOPSD outperform standard on-policy self-distillation (OPSD)?

Standard OPSD relies on external reference solutions as privileged information, which the student cannot access at inference, leading to a weak "consensus" policy. dOPSD sources privilege from the model's own denoising trajectory, which is intrinsic to the decoding process and available without external labels.

Does dOPSD require ground-truth solutions to function?

No. While the authors use rollout verification against gold answers to discard incorrect samples, they demonstrate that dOPSD still improves over the base model even without this verification, making it applicable to datasets lacking reference answers.

Researchers can now treat the denoising trajectory of diffusion models as a source of privileged information, enabling label-efficient post-training that avoids the pitfalls of external teacher dependence.

Introduction

We expose why standard OPSD fails for diffusion models and outline our dOPSD solution.

Standard on‑policy self‑distillation (OPSD) assumes a teacher can look at a ground‑truth reference that the student never sees at inference; this mismatch collapses to a weak, PI‑free consensus policy and yields little reasoning gain for diffusion large language models (dLLMs).

OPSD’s teacher is privileged by an external solution that the model cannot access when generating text, so the student ends up learning from a diluted signal that does not improve reasoning.

dLLMs need better post‑training methods to unlock strong reasoning abilities.

Background: Diffusion Language Models

Survey of diffusion language models, distillation paradigms, and privileged self‑teaching approaches.

Related work clusters around three themes: diffusion language modeling, distillation techniques, and privileged‑information self‑teaching.

A dLLM generates text by iteratively denoising a partially masked sequence, rather than emitting tokens left‑to‑right.

Diffusion models have been scaled to billions of parameters (e.g., LLaDA and Dream) while preserving parallel, confidence‑ordered decoding.

Models such as LLaDA and Dream extend the diffusion paradigm to massive parameter counts, keeping the mask‑based denoising loop and enabling bidirectional context.

Reinforcement‑learning‑with‑value‑regularization (RLVR) is adapted to diffusion models by introducing surrogate likelihoods or in‑painting guidance because diffusion lacks a tractable sequence log‑likelihood.

Transfers softened output probabilities from a large teacher model to a smaller student, improving generalization.

Trains a student on its own generations, using a separate, larger teacher to provide dense feedback along the rollout.

Uses a teacher conditioned on extra information unavailable at test time (e.g., a ground‑truth solution or an in‑context prompt) to guide the student along its own rollouts.

dOPSD unifies these strands: it applies on‑policy self‑distillation to diffusion language models, using the model’s own later denoising state as the privileged teacher, thereby delivering dense token‑level supervision without an external teacher.

The OPSD Framework

OPSD lets a model teach itself by comparing a no‑context rollout to a privileged‑context rollout.

Standard OPSD assumes an autoregressive token‑by‑token teacher, which does not line up with the iterative denoising trajectories of diffusion language models.

OPSD lets the model grade its own generated completions by comparing a “student” rollout that sees only the prompt to a “teacher” rollout that also sees the reference answer – like a student checking its essay against an answer key it wrote earlier.

Student draws $\hat{y}_{1}\sim p_{S}(\cdot\mid x)$ → selects token “a”.

Teacher, conditioned on $x$ and $y^{\star}$, computes $p_{T}(\cdot\mid x, y^{\star}, \hat{y}_{<1})$ and assigns high probability to “yes”.

Compute divergence $D\big(p_{T}(\cdot\mid x, y^{\star}, \hat{y}_{<1})\;\|\;p_{S}(\cdot\mid x, \hat{y}_{<1})\big)$ and back‑propagate into the student.

Student draws $\hat{y}_{2}\sim p_{S}(\cdot\mid x, \hat{y}_{1})$ → selects token “b”.

Teacher evaluates the second position with the same privileged context and yields a second divergence term.

Average the two divergences; the student’s parameters are updated to reduce the gap.

OPSD provides a token‑wise learning signal even when the reference answer is only available during training, eliminating the need for an external teacher model.

How does OPSD differ from classic self‑distillation that uses a fixed external teacher?

In classic self‑distillation the teacher is a separate, pre‑trained model that never sees the ground‑truth answer. OPSD re‑uses the same parameters but gives the teacher privileged access to $y^{\star}$, so the teacher’s distribution is conditioned on information the student cannot use at inference, yielding on‑policy, token‑level targets.

The privileged teacher is a copy of the model that is allowed to peek at the reference solution during training, giving it “extra eyes” that the student never gets at test time.

Why call it a “privileged” teacher instead of simply a stronger model?

The term highlights that the teacher’s advantage comes solely from additional conditioning information ($y^{\star}$), not from extra parameters or training data. At inference the model loses that privilege, so the student must learn to emulate the teacher’s behavior without the extra context.

Why OPSD Fails for dLLMs

OPSD breaks for diffusion models because its teacher collapses and its masking drifts off the decoding path.

A diffusion language model walks a denoising path: starting from a fully masked sequence it repeatedly fills in the most confident tokens, freezing them for the rest of the generation.

Step 1: model predicts token “A” with highest confidence at position 2 and commits it → $s_1=[\langle\text{mask}\rangle, A, \langle\text{mask}\rangle,\langle\text{mask}\rangle]$.

Step 2: next highest confidence is “B” at position 4 → $s_2=[\langle\text{mask}\rangle, A, \langle\text{mask}\rangle, B]$.

Step 3: “C” at position 1 → $s_3=[C, A, \langle\text{mask}\rangle, B]$.

Step 4: final token “D” at position 3 → $s_4=[C, A, D, B]=\hat y$.

The trajectory shows a confidence‑ordered, monotonic mask shrinkage; any state outside this order cannot be visited by the model.

Applying OPSD to a diffusion model encounters two fundamental obstacles. First, the privileged‑teacher formulation collapses into a PI‑free policy that averages over instance‑specific references, eroding the benefit of a teacher that knows the exact solution. Second, the random‑mask supervision drifts far from the model’s own decoding frontier, feeding the student signals it never sees at inference.

When the teacher is conditioned on a distribution of reference solutions instead of the single ground‑truth answer, the student is forced to match the geometric mean of many divergent targets, which weakens any instance‑specific signal.

How does this collapse differ from the usual teacher‑student distillation where the teacher is fixed?

In standard distillation the teacher provides a single deterministic output, so the student learns that exact mapping. Here the “teacher” is a mixture over many possible reference solutions, and the student must satisfy the average of their log‑probabilities, which blurs the signal and removes instance‑specific guidance.

Randomly masking tokens of a completed rollout creates partial states that the diffusion model never visits, because during generation the mask follows a confidence‑ordered schedule rather than a uniform random pattern.

Why can’t we simply mask random tokens of the final rollout and still train the student?

Because the model’s inference path never visits those random‑masked states; the student would be learning to predict tokens under a context that never occurs, leading to a mismatch between training and test‑time behavior.

**Figure 2.** OPSD cannot be applied directly to diffusion language models (Dream-7B-Instruct). (a) The fraction of a uniformly masked subset that lands above the on-policy decoding frontier, easy tokens the model would already have committed, is large at small mask fractions p, whereas a mask taken from a genuine decoding step (blue) is zero by construction. (b) Scoring the student at random mask tokens of the finished rollout (off-path) degrades GSM8K accuracy over training, while scoring it at genuine random decode steps stays high. (c) Ported directly to a dLLM, both OPSD variants train to substantially lower accuracy than our trajectory-based dOPSD, and conditioning the teacher on the full reference solution collapses performance.

The dOPSD Method

Distill a model from its own decoding trajectory by averaging future predictions at masked positions.

Standard OPSD cannot exploit the denoising path of diffusion language models because it relies on token‑wise autoregressive probabilities that never appear during inference.

The model reuses its own decoding trajectory so that the student sees a genuine intermediate state while the teacher looks ahead along the same trajectory, providing privileged context without any external reference.

For masked position $i=2$, future steps where it stays masked are $T_2=\{2,3\}$ because $i$ is still masked in $s_2$ and $s_3$.

Compute teacher predictions $\pi_\theta(\,\cdot\mid x,s_2)_2$ and $\pi_\theta(\,\cdot\mid x,s_3)_2$, then average: $\bar p_2=\tfrac12\bigl(\pi_\theta(\,\cdot\mid x,s_2)_2+\pi_\theta(\,\cdot\mid x,s_3)_2\bigr)$.

For masked position $i=3$, $T_3=\{3\}$ (only the final state keeps it masked), so $\bar p_3=\pi_\theta(\,\cdot\mid x,s_3)_3$.

The student at step $k=1$ predicts $\pi_\theta(\,\cdot\mid x,s_1)_2$ and $\pi_\theta(\,\cdot\mid x,s_1)_3$, which are then compared to $\bar p_2$ and $\bar p_3$ via the JSD loss.

Because the teacher predictions incorporate more surrounding tokens, the divergence pushes the student toward a better‑informed distribution.

Averaging over all remaining masked steps guarantees that the teacher never sees a token it has already generated, avoiding the “copy‑the‑answer” shortcut that would otherwise inflate the loss.

**Figure 1.** Overview of dOPSD. (a) A single shared model plays both the student and the teacher. Given a problem $x$ with reference answer $y^*$ sampled from the dataset, the model rolls out an on-policy decoding trajectory $\xi = s_i arrow \dots arrow s_K arrow \hat{y}$. The student is scored at an intermediate, still-masked step $s_i$, $p_s = \pi_\theta(. | x, s_i)$, while the same model acting as the teacher scores the same positions from later, more-decoded states of the trajectory, an on-policy peek-ahead that uses no external solution. A rollout verifier $\beta = 1[\hat{y} = y^*]$ keeps only correct rollouts, and the student is trained to match the stop-gradient teacher target, with gradients flowing only into the student. (b) Construction of that teacher target at a masked position: the teacher's predictive distribution is averaged over the future decoding steps where the position is still masked (solid, the set $\mathcal{T}_i$) and dropped once the position has been decoded (hatched, where the model would merely copy the committed token); the average is the teacher target $\bar{p}_i$ the student is distilled toward.

Algorithm 1: dOPSD training loop

Experimental Setup

Key experimental conditions for dOPSD and its baselines.

dOPSD is evaluated alongside four post‑training baselines and the untuned base model.

We compare dOPSD against four post‑training methods trained on the same data and base models, together with the untuned base model.

All methods are fine‑tuned with parameter‑efficient LoRA (rank 32, $\alpha = 32$) applied to the query, key, value, output, gate, up‑ and down‑projection matrices. Training runs on 8 GPUs with a per‑device mini‑batch size of 4, using AdamW at a peak learning rate of $2 \times 10^{-5}$ for 3 epochs; we report the best checkpoint for each method. For dOPSD the trajectory mask threshold is set to $\tau = 0.5$, ensuring the student is scored when at least half of the answer region remains masked.

Evaluation measures in‑domain mathematical reasoning on GSM8K and MATH500 (accuracy) and out‑of‑distribution code generation on HumanEval and MBPP (pass@1), using the trained checkpoints without any further fine‑tuning.

Main Results

dOPSD outperforms all baselines on every benchmark, both in‑domain and OOD.

dOPSD improves code generation on Dream‑7B‑Instruct by $+4.17$ points over the base model.

HumanEval score rises from 52.54 to 56.71 (Table 1).

**Table 1.** Performance of dOPSD against supervised, reinforcement-learning, and on-policy self-distillation baselines on Dream-7B-Instruct and LLaDA-8B-Instruct. All methods are trained on the mathematical-reasoning corpus; code generation is evaluated out of distribution. The best result in each column is in **bold** and the second best is <u>underlined</u>.

Ablation Studies

Ablations quantify how each design choice—KL direction, teacher horizon, and mask threshold—affects performance.

We isolate three components of dOPSD to see how each contributes to the final scores: the KL divergence direction, the length of the teacher’s look‑ahead window, and the mask‑threshold $\tau$ that controls which parts of the trajectory are used.

Using forward KL ($\beta$→0) instead of reverse KL ($\beta$→1) raises Dream MBPP from 50.61 to 58.49.

Figure 3 shows forward KL beating reverse KL on every benchmark for both Dream‑7B‑Instruct and LLaDA‑8B‑Instruct.

Figure 3 visualizes the two KL extremes; forward KL consistently yields higher scores because it forces the student to cover the full teacher distribution rather than collapsing onto a single mode.

Extending the teacher horizon to the full remaining trajectory improves GSM8K from 78.14 (200 steps) to 83.04.

Table 2 reports performance for horizons of 0, 50, 100, 200, and the full trajectory, with the full horizon dominating all columns.

Table 2 demonstrates that a longer look‑ahead supplies richer privileged signals, turning the teacher into a more informative guide for the student.

Setting the mask threshold $\tau$ to 0.50 yields the best average score (60.11), a modest gain of +0.6 over the nearest alternatives.

Table 3 compares $\tau$ = 0.25, 0.50, 0.75; the middle value consistently achieves the highest numbers across GSM8K, MATH500, HumanEval, and MBPP.

Table 3 confirms that a balanced mask—neither too permissive nor too restrictive—captures enough of the decoding trajectory to maximize the distilled signal.

Qualitative Analysis

Qualitative examples show dOPSD avoids baseline errors on representative GSM8K problems.

dOPSD solves both qualitative examples correctly while every baseline fails.

On the two GSM8K problems, dOPSD produces the gold answer (21 and 2.5) whereas all baselines generate incorrect or incomplete results.

The image displays a math problem and its solution generated by a model named dOPSD. The problem asks for the weight difference after 5 weeks between two breakfast scenarios (losing 1.25 lbs/week vs. gaining 1.75 lbs/week). The dOPSD model calculates the total loss (6.25 lbs) and total gain (8.75 lbs), then sums them to arrive at the correct answer of 15.

The image displays three distinct model output examples labeled GRPO, OPSD (answer-only), and OPSD (full-solution), each with a corresponding predicted value. Below these examples, a paragraph describes the failure modes of the baselines, noting that they make confident slips or produce malformed/empty generations, while dOPSD maintains reasoning faithfulness.

Additional Qualitative Examples

Qualitative examples show dOPSD consistently hits the gold answer while baselines miss.

Across the two qualitative examples the dOPSD method uniquely preserves the signed gap and then reports the positive magnitude, whereas every baseline either drops the sign or collapses the result.

dOPSD recovers the gold answer on both qualitative examples.

Example 2 (weight‑difference) and Example 3 (score‑difference) both match the gold answers of 15 and 25 respectively.

Baselines consistently make a single slip: they either drop the sign, treat the signed gap as the final answer, or collapse ambiguous terms to zero, causing a cascade of errors.

Conclusion

Conclusion summarizes dOPSD’s advantages, empirical gains, and practical implications.

We introduced dOPSD, an on‑policy self‑distillation technique that reinterprets the privileged‑teacher principle for diffusion language models.

Standard OPSD assumes an external, instance‑specific reference solution as privileged information; this reference disappears at inference, causing OPSD to collapse onto a weak, reference‑free consensus, and its random masking misaligns with the model’s confidence‑ordered decoding trajectory.

dOPSD resolves both problems by sourcing the teacher’s privilege and the student’s noise from a single origin—the student’s own denoising trajectory—so the student is evaluated at a genuine intermediate step while the same model, acting as teacher, evaluates the identical positions from later, more‑decoded steps, providing an honest peek‑ahead without any external reference.

Empirically, on Dream‑7B‑Instruct and LLaDA‑8B‑Instruct, dOPSD is the only method that consistently improves the base model, boosting in‑domain mathematical reasoning and, despite training solely on mathematics, also enhancing out‑of‑distribution code generation, whereas supervised, RL, and naive OPSD baselines either stagnate or collapse.

Ablation studies confirm that the gains originate from the trajectory‑derived privileged signal itself; they persist even when rollout verification is omitted, demonstrating that dOPSD can be applied to datasets lacking reference answers.

Table 2 examines the teacher horizon, i.e., how many future trajectory steps are averaged into the teacher target $\bar{\pi}$ (Eq. 6); the “Full” setting, which averages over the entire remaining masked trajectory, yields the best performance.

Table 3 varies the trajectory mask threshold $\tau$, which controls how masked the selected student step must be; a moderate $\tau = 0.5$ achieves the highest average (60.11) and the top score on every benchmark except HumanEval.

When rollout verification is removed, dOPSD still raises the four‑task average from 57.59 to 58.64 and improves GSM8K, MATH500, and HumanEval, with only a slight drop on MBPP; verification adds a further boost to 60.11 but is not essential to the method’s success.

Dataset Details

Details of the training corpus and evaluation benchmarks used in experiments.

We post‑train every method on MixChain‑Z‑PRM12K, a ≈12 K problem corpus derived from the MATH/PRM800K lineage. Each record contains a checkable final answer and one or more chain‑of‑thought (CoT) solutions; the mixed‑chain construction supplies both terse derivations and fully elaborated step‑by‑step traces, exposing the model to a spectrum of reasoning granularities. The final answer feeds the rollout verifier of $Eq.~(8)$, while reference CoT solutions are only consumed by baselines such as SFT and full‑solution OPSD.

dOPSD itself consumes neither the reference CoT nor, in the unverified ablation of Figure 4, the final answer. All methods use the same instruction‑style prompt template as at evaluation time, tokenized with each backbone’s native tokenizer and truncated to a 256‑token completion budget that matches the diffusion decoding length. Problems whose gold answer cannot be parsed into a checkable form are discarded.

Evaluation uses four public benchmarks: GSM8K (1 319 test items) and MATH500 (500 items) for in‑domain mathematical reasoning, scored by final‑answer accuracy (Acc@1); HumanEval (164 items) and MBPP (500 items) for out‑of‑distribution code generation, scored by functional correctness (pass@1). For all benchmarks we decode with the same diffusion sampler and generation length as during training, reading out the final answer greedily, so reported numbers reflect single‑sample accuracy.

**Table 4.** Evaluation benchmarks. Math benchmarks are scored by final-answer accuracy (Acc@1); code benchmarks by functional correctness against the provided unit tests (pass@1). MBPP uses the sanitized test split.

Inference Hyperparameters

Hyperparameter settings used for inference across models and benchmarks.

Table 5 presents the full decoding configuration for every evaluation, broken down by model backbone and benchmark. The table is the single source of truth for the hyperparameters used in inference.

All numbers in the paper are obtained with single‑sample decoding; consequently, math benchmarks report Acc@1 while code benchmarks report pass@1. No best‑of‑n or self‑consistency tricks are employed.

Two invariants hold across every run. First, the number of diffusion steps is forced to equal the number of generated tokens, because committing more than one token per step sharply harms quality.

Second, the generation length is chosen per benchmark to comfortably cover the reference solutions while staying within each model’s position limit; this is why code benchmarks allocate more tokens than the math benchmarks.

Dream‑7B‑Instruct decodes with confidence‑based (entropy) remapping at a low temperature of 0.1 and nucleus sampling with top‑p = 0.9, without any top‑k truncation. For the HumanEval code benchmark it additionally enables entropy‑threshold parallel decoding with a threshold of 0.5 for efficiency.

LLaDA‑8B‑Instruct uses its native low‑confidence remapping at temperature 0 (greedy). Its Gumbel‑argmax sampler ignores nucleus truncation, so the listed top‑p = 0.9 is purely informational. The model also denoises in semi‑autoregressive blocks, with block length set per benchmark (8 on GSM8K up to a full‑length block on HumanEval).

**Table 5.** Inference hyperparameters across base models and benchmarks. “–” indicates the hyperparameter is not used. On both backbones the number of diffusion steps equals the number of generated tokens. Temperature is 0 (greedy) and top-$p$ is informational only, as its Gumbel-argmax sampling ignores nucleus truncation.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers