PaperlayerView exact revision on arXiv

Source-linked AI summary

d-OPD: Future-Aware On-Policy Distillation for Block Diffusion Language Models

Ruitao Liu, Qinghao Hu, Song Han

arXiv:2609.35362v1cs.LGcs.AI

TL;DR

In AR-to-block-diffusion distillation, the causal teacher conditions only on the prefix while the student can use visible future context in its block. d-OPD uses that context to correct teacher supervision, and across Qwen3 models it improves benchmark averages over OPDLM while reducing training time.

  • Problem

    The causal AR teacher target uses only the prefix, whereas the block-diffusion student can condition on the full visible block, including future context.

    Source: §1 · p. 2, §3.1 · p. 4

  • Method

    d-OPD uses visible future context to correct the AR teacher distribution, approximating the target conditioned on the student-visible state from completed on-policy trajectories.

    Source: §1 · p. 2

  • Results

    Across Qwen3 models from 0.6B to 8B, d-OPD gains 2.6–4.0 points in the six-benchmark average over OPDLM.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Across the evaluated Qwen3 scales, correcting the teacher target improves the six-benchmark average and reduces time to matched OPDLM performance.

    Source: §6 · p. 9

    Computing the exact future-conditioned target is intractable because it requires marginalizing over compatible left-prefix completions and unobserved future tokens.

    Source: §4.1 · p. 6

Abstract

from arXiv · show

Large language models (LLMs) typically generate text autoregressively (AR), predicting one token at a time. Block diffusion language models (dLLMs) instead generate blocks sequentially while denoising multiple tokens in parallel within each block, offering a promising way to accelerate generation. Rather than training such models from scratch, recent work adapts strong pretrained AR models into block dLLMs through distillation. On-policy distillation (OPD) has been widely used for LLM training because it supervises the student on states generated by its current policy, rather than only on fixed offline trajectories. By training on the states the student actually visits, it reduces the mismatch between training and generation and can provide more relevant supervision as the student evolves. Recent work has extended this idea to AR-to-block-diffusion conversion. However, this setting introduces a fundamental mismatch in supervision: the block-diffusion student and the causal AR teacher condition on different information at the same training state. The student predicts from the entire partially denoised block, including visible future context, whereas the standard AR teacher target is defined only from the causal prefix. As a result, the teacher distribution used for distillation is not fully aligned with the information available to the student. We therefore introduce d-OPD, a future-aware on-policy distillation method that corrects the AR teacher distribution to better align with the student-visible state by incorporating visible future information within each block, providing supervision that better matches the information used by the student. Across Qwen3 models from 0.6B to 8B, d-OPD improves the six-benchmark average by up to $4.0$ points over OPDLM and reduces training time by $1.35$-$1.58\times$. The code is available at https://github.com/mit-han-lab/d-OPD.

1 INTRODUCTION

Block diffusion enables parallel token generation within causal blocks, but AR-to-block-dLLM distillation retains a mismatch between the information available to the student and the causal teacher target. d-OPD uses visible future context to correct that target, improving benchmark performance and training efficiency across evaluated Qwen3 models.

1 INTRODUCTION

  • Block dLLMs generate blocks from left to right while allowing tokens within the current block to interact bidirectionally.This permits parallel generation within a block while preserving causal structure across blocks.Source: §1 · p. 1
  • The block-diffusion student can use visible tokens to the right of a masked position, while the causal AR teacher target depends only on the prefix.Thus, querying the teacher on student-visited states does not by itself align the information used by teacher and student.Source: §1 · p. 2
  • Visible future tokens can provide evidence about a masked token because candidate values make the observed future more or less likely under the AR teacher.Source: §1 · p. 2
  • d-OPD reweights the causal AR teacher distribution using visible future context, approximating the target conditioned on the complete student-visible state from completed on-policy trajectories.Exact conditioning requires intractable marginalization over undenoised prefix and future tokens.Source: §1 · p. 2
  • Up to 4.0 points higher six-benchmark average and 1.35–1.58× faster time to matched OPDLM performance across Qwen3 models from 0.6B to 8B.Figure 1 also reports a 71.9% reduction in mean token-level KL after future-aware correction.Source: §1 · p. 2, Figure 1 · p. 3

2 RELATED WORK

Related work covers diffusion-language-model inference, conversion of pretrained AR models to diffusion models, and on-policy distillation. The paper distinguishes its focus on the teacher target used during conversion from prior inference improvements and OPD's alignment of student states.

2.1 DIFFUSION LANGUAGE MODELS

  • Block dLLM research includes inference improvements such as confidence-aware decoding, cache reuse, block schedules, generation orders, and serving systems.These works mainly address inference for already-trained diffusion language models.Source: §2.1 · p. 3

2.2 AR-TO-DIFFUSION ADAPTATION

  • AR-to-diffusion conversion methods adapt pretrained models through changes to attention, masking, denoising, and block-wise training while seeking to preserve pretrained knowledge.BARD, T⋆, and related work also use progressive scaling, stage-wise distillation, curricula, auxiliary AR objectives, and on-policy adaptation.Source: §2.2 · p. 3

2.3 ON-POLICY DISTILLATION

  • On-policy distillation trains on student-generated states, and OPDLM applies it to AR-to-dLLM conversion but leaves a target-conditioning mismatch.The remaining mismatch is between the causal teacher and block-dLLM student conditioning on different information.Source: §2.3 · p. 3

3 MOTIVATION

Block-dLLM students can use visible tokens to the right of a masked position, while the causal AR teacher target uses only the prefix. d-OPD conditions teacher supervision on the student-visible state, and a controlled diagnosis finds that its target is closer to the exact posterior.

3.1 TARGET-CONDITIONING MISMATCH

  • The target-conditioning mismatch arises because a block dLLM can condition on visible tokens to the right of a masked position, whereas the causal AR target uses only the prefix.Thus the student and teacher targets depend on different information at the same student-visited state.Source: §3.1 · p. 4

3.2 EXACT FUTURE-CONDITIONED TARGET

  • The block-level target conditions the teacher joint distribution on completions consistent with the visible block, then takes the token-level marginal at the masked position.The exact target accounts for uncertainty in masked prefixes and unobserved futures; with no visible future, it reduces to the causal prior.Source: §3.2 · pp. 4–5
  • Figure 2 traces how d-OPD uses a completed on-policy block to score candidate tokens against the visible future, reweights the causal prior, and trains the student by KL divergence.The completed prefix supplies the causal AR teacher prior.Source: Figure 2 · p. 5
  • The exact token-level teacher target reweights the causal prior by how likely the visible future is under each candidate token.This yields the teacher posterior conditioned on the complete student-visible state.Source: §3.2 · pp. 4–5

3.3 DIAGNOSIS

  • 0.4444 mean token-level KL for the causal target versus 0.1249 for d-OPD, a 71.9% reduction in distance to the exact posterior across 49,948 masked-token marginals.The controlled setting computes the exact posterior by enumeration.Source: §3.3 · p. 5

4 METHOD

d-OPD approximates a teacher target conditioned on the full student-visible block by using the completed on-policy rollout. It applies this future-aware correction selectively, while retaining the existing student and training procedure.

4.1 FUTURE-AWARE CORRECTION

  • d-OPD approximates the future-conditioned teacher target using the rollout’s completed prefix and suffix, since exact computation requires intractable marginalization.The completed suffix likelihood under each candidate token estimates future compatibility.Source: §4.1 · p. 6
  • To reduce vocabulary-wide scoring cost, future scores are computed only for a small set of high-probability candidates and normalized relative to the rollout token, giving scored candidates a neutral baseline while unscored tokens retain correction weight one.The candidate set is the union of teacher and student top-k sets; top-4 tokens cover more than 95% of probability mass for both distributions.Source: §4.1 · p. 6, §4.2 · p. 7

4.2 METHOD DETAILS

  • Future-aware correction is gated on verifier-correct rollouts; incorrect rollouts still receive standard causal teacher supervision.The gate avoids future scoring on incorrect rollouts, whose futures may provide noisy or misleading reweighting evidence.Source: §4.2 · p. 7

4.3 TRAINING

  • Each block contributes one randomly sampled on-policy denoising state, and distillation applies forward KL at every masked position in that state.Token losses are averaged across masked positions; d-OPD changes teacher supervision while keeping the student architecture, objective, and overall training procedure unchanged.Source: §4.3 · pp. 7–8

5 EXPERIMENTS

Experiments compare d-OPD with OPDLM and other conversion baselines across Qwen3 scales and block sizes. d-OPD improves benchmark averages and reaches matched OPDLM performance faster, while ablations examine candidate restriction and correctness gating.

5 EXPERIMENTS

  • Experiments use Qwen3 checkpoints from 0.6B to 8B, with block size N = 4 for scale comparisons and additional Qwen3-8B runs at N = 8 and 16.The frozen pretrained AR checkpoints serve as both causal teacher and future-compatibility scorer.Source: §5 · p. 8

5.1 MAIN RESULTS

  • d-OPD gains 2.6–3.9 points in six-benchmark average over OPDLM across Qwen3 scales, and has the strongest converted-model average against other baselines on Qwen3-8B and Qwen3-4B.At larger block sizes, gains reach 4.0 points; Table 1 reports accuracy across six benchmarks, model scales, and block sizes.Source: §5.1 · p. 8

5.2 TRAINING EFFICIENCY

  • d-OPD reaches matched OPDLM performance 1.35–1.58× faster across the four Qwen3 scales, measured by wall-clock time to the OPDLM best six-benchmark average.The timing comparison accounts for the additional teacher computation from future scoring.Source: §5.2 · p. 9

5.3 ABLATION STUDIES

  • At k = 16, candidates cover more than 99% of teacher and student probability mass with 2.63% training overhead; larger sets yield marginal coverage gains but higher overhead.The authors use k = 16 as the default coverage–cost tradeoff.Source: §5.3 · p. 9
  • Correctness gating raises the six-benchmark average from 58.8 to 59.8 by applying future-aware correction only to verifier-correct rollouts.The gate also avoids unnecessary future scoring on incorrect rollouts.Source: §5.3 · p. 9, Table 3 · p. 9

6 CONCLUSION

The paper concludes that d-OPD improves OPDLM results by using student-visible future context to correct the teacher target. It also documents AI use, ethics, and reproducibility procedures.

6 CONCLUSION

  • Across Qwen3 models from 0.6B to 8B, d-OPD improves over OPDLM by up to 4.0 points in six-benchmark average and reaches matched performance up to 1.58× faster.The method reweights the causal AR teacher distribution using visible future context within each block.Source: §6 · p. 9

AI USE STATEMENT

  • The authors report using generative AI for implementation assistance, language polishing, evaluation and experiment code, and execution monitoring, with reviewed AI-assisted work.They state AI was not used for research ideation, methodology or experiment design, hypothesis development, mathematical claims or proofs, data analysis, or result interpretation.Source: AI USE STATEMENT · p. 10

ETHICS STATEMENT

  • The work uses existing models, benchmarks, and training data, and reports no human subjects or private or personally identifiable information collection.The authors report no additional method-specific ethical concerns beyond broader considerations of training and deploying LLMs.Source: ETHICS STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

  • The paper documents algorithm, training, evaluation, diagnosis, and timing protocols across its main text and appendices to support reproducing the reported comparisons.It also documents implementation details, including candidate scoring and caching, and full correctness-gating ablations.Source: REPRODUCIBILITY STATEMENT · p. 10

A TRAINING ALGORITHM

The algorithm constructs future-aware teacher targets from student-visited states and completed trajectories during on-policy training.

A TRAINING ALGORITHM

  • For each masked prediction in student-visited denoising states, d-OPD builds a future-aware teacher target from the same rollout’s completed trajectory.Correction is limited to a small candidate set and verifier-passing rollouts; otherwise, training uses the ordinary causal teacher target.Source: A TRAINING ALGORITHM · p. 17

B ADDITIONAL IMPLEMENTATION NOTES

d-OPD reuses cached AR prefix computation to score candidate continuations and correct teacher targets on student-visited states. Training uses the OPDLM prompt mixture, with future-aware correction gated by response correctness.

B ADDITIONAL IMPLEMENTATION NOTES

  • Candidate scoring reuses cached causal-prefix computations across branches, scoring the union of teacher and student top-k candidates and limiting cache expansion to 16 GiB.The default k = 16 job uses at most 33 branches, including a reference branch; caches are reused within a trajectory but not across examples.Source: B ADDITIONAL IMPLEMENTATION NOTES · pp. 17–18
  • For each sampled visited block state, d-OPD forms a causal teacher prior and, when eligible, reweights candidate tokens by relative suffix likelihood before normalizing the corrected target.The correction is applied when the rollout is verifier-correct and visible future context exists; otherwise the causal target is retained.Source: B ADDITIONAL IMPLEMENTATION NOTES · p. 18

C.1 DATA

  • Training uses the 61,816-prompt OPDLM mixture with inherited preprocessing and Qwen chat formatting, without additional deduplication or filtering.Source: §C.1 · pp. 18–19

C.2 TRAINING

  • Each comparison uses the same prompt mixture, rollout configuration, state selection, optimization budget, and seeds, with one pass over 483 on-policy iterations.Training runs on eight H100 80 GB GPUs; the default candidate size is k = 16.Source: §C.2 · p. 19
  • Training reconstructs denoising states visited during rollout, samples a masked-prediction round per block, and keeps each rollout as a separate trajectory-level example.Up to 128 prompts are processed per iteration; correctness gating enables future-aware correction only for verifier-correct rollouts, while incorrect or unverifiable examples retain causal supervision.Source: §C.2 · p. 19, §C.3 · p. 19

C.4 BENCHMARK EVALUATION

Benchmark evaluation uses zero-shot Qwen3 completions and selects checkpoints by their six-benchmark average. A separate controlled diagnostic constructs an enumerable teacher support to compare approximate targets with the exact future-conditioned posterior.

C.4 BENCHMARK EVALUATION

  • The six benchmarks are evaluated zero-shot, and all reported benchmark scores come from the checkpoint with the best unweighted six-benchmark average.The evaluation uses the full benchmark sets and corresponding block length, four denoising steps per block, and one completion per problem.Source: §C.4 · pp. 19–20

D.1 RESTRICTED ENUMERABLE SUPPORT

  • The diagnostic constructs a restricted enumerable vocabulary sequentially, retaining each position’s top-V tokens under a teacher-weighted mixture of restricted prefixes.This creates a finite support intended to preserve high-probability local choices while allowing exhaustive enumeration.Source: D ENUMERABLE TEACHER CONSTRUCTION AND… · p. 20, §D.1 · p. 20

D.2 EXACT TEACHER POSTERIOR

  • The exact reference posterior is obtained by scoring every complete sequence under the original AR factorization and renormalizing over the restricted support.For each partially observed state, enumeration yields the exact block posterior and token marginal against which approximate targets can be compared.Source: §D.2 · p. 20

D.3 COMPARED TARGETS

  • The future-aware target reweights the causal teacher distribution using each candidate’s compatibility with the same completed future, then normalizes over the restricted support.Unlike large-scale training, this diagnostic applies the correction across the full restricted vocabulary rather than using a top-k approximation.Source: §D.3 · pp. 20–21

D.4 KL METRIC AND AGGREGATION

The diagnostic measures token-target quality against an exactly enumerated future-conditioned target using KL divergence. Its controlled sweep varies block length, masking, and student-distribution perturbation, then pools token-level values across the resulting states.

D.4 KL METRIC AND AGGREGATION

  • Target quality is the KL divergence between the exact future-conditioned token target and either the causal or future-aware estimate, computed over the restricted candidate support.Source: §D.4 · p. 21
  • The diagnostic tests block lengths N ∈ {4, 6, 8} using 400 sampled blocks per configuration, with positions independently retained at probability 0.45 and at least one masked target per state.Source: §D.4 · p. 21
  • Student-distribution perturbation levels σ ∈ {0, 0.5, 1.0, 1.5} range from the restricted teacher distribution to increasingly perturbed distributions, with the same observed states reused.A compatible completion is sampled at each noise level to construct the causal and future-aware targets.Source: §D.4 · p. 21
  • When no future token is visible to a masked position’s right, the future-compatibility term is constant across candidates and the two targets coincide.The aggregate diagnostic is the arithmetic mean of token-level KL values pooled over 49,948 masked-token marginals.Source: §D.4 · p. 21

E WALL-CLOCK TIMING PROTOCOL

With the default k = 16 configuration, d-OPD adds 2.63% average per-iteration overhead over OPDLM. The timing protocol also compares time to a matched benchmark threshold, while the paper identifies reliance on a single completed trajectory as a limitation.

E WALL-CLOCK TIMING PROTOCOL

  • Matched-performance timing compares when each method reaches OPDLM’s best six-benchmark average, while d-OPD timing includes routine training and checkpointing costs.Source: E WALL-CLOCK TIMING PROTOCOL · pp. 21–22
  • 2.63% is d-OPD’s average per-iteration overhead over OPDLM across four model scales under the default k = 16 configuration.The timed iteration includes rollout generation, verification, teacher scoring, student optimization, serialization, and weight reloading; periodic benchmark evaluation is excluded.Source: E WALL-CLOCK TIMING PROTOCOL · p. 22

F LIMITATIONS AND FUTURE WORK

  • Current d-OPD constructs future-aware targets from a single completed on-policy trajectory; aggregating multiple compatible trajectories remains proposed future work.The authors suggest multiple futures may reduce variance, and identify evaluation across additional model families and block sizes as future work.Source: F LIMITATIONS AND FUTURE WORK · p. 22

G ADDITIONAL ABLATION RESULTS

The correctness-gating ablation reports full results for Qwen3-8B. Its gate applies future-aware correction only to verifier-correct rollouts, while ordinary causal teacher supervision remains active for all rollouts.

G ADDITIONAL ABLATION RESULTS

  • Table 5 reports the full per-benchmark results for the correctness-gating ablation on Qwen3-8B.The gate applies future-aware correction only to verifier-correct rollouts; ordinary causal teacher supervision remains active for all rollouts.Source: G ADDITIONAL ABLATION RESULTS · p. 22, Table 5 · p. 22
  • The gate restricts future-aware correction to verifier-correct rollouts, while causal teacher supervision remains active for every rollout.Source: Table 5 · p. 22
  • The reported Avg. is the unweighted average over six benchmarks, and higher scores are better.Source: Table 5 · p. 22
Loading 2609.35362v1…