PaperlayerView exact revision on arXiv

Source-linked AI summary

Data-Efficient Autoregressive-to-Diffusion Language Models via On-Policy Distillation

Xingyu Su, Jacob Helwig, Shubham Parashar, Atharv Chagi, Lakshmi Jotsna, Degui Zhi, James Caverlee, Dileep Kalathil, Shuiwang Ji

arXiv:2606.06712v1cs.CLcs.AI

TL;DR

Converting autoregressive models into diffusion language models can lose learned knowledge and mismatch inference states. OPDLM uses on-policy distillation from the original model and achieves competitive performance across tasks with 15× to 7,000× fewer training tokens.

  • Problem

    ARLM-to-DLM conversion faces knowledge-retention and training–inference state mismatches between objective changes, random masking, and reverse decoding trajectories.

  • Method

    OPDLM initializes a bidirectional student from the ARLM, samples its reverse trajectories, and distills token distributions from the frozen ARLM teacher on those states.

  • Results

    15× to 7,000× fewer training tokens yields competitive performance across a broad range of tasks, with weaker coding results where corpus coverage is limited.

  • Takeaways & Limitations

    The results support treating ARLM-to-DLM conversion as an efficient post-training procedure rather than computationally intensive DLM pretraining.

  • Takeaways & Limitations

    OPDLM remains weaker than the strongest baselines on certain tasks, while controlled comparisons are difficult because baseline pretraining datasets are unavailable.

Abstract

from arXiv · show

We study the transformation of autoregressive models (ARLMs) into diffusion language models (DLMs). Rather than pretraining from scratch, prior work replaces the causal attention in ARLMs with bidirectional attention and then trains the resulting model using a DLM objective. However, these approaches incur two distribution shifts. First, transitioning from a next-token prediction objective to a DLM objective can discard knowledge acquired by the ARLM during training. Second, standard DLMs suffer from a train-inference mismatch, as the training loss is defined on randomly masked sequences rather than the trajectories encountered at inference produced by confidence-based decoding. To address both challenges, we introduce an On-Policy Diffusion Language Model (OPDLM) in which On-Policy Distillation (OPD) is employed for ARLM-to-DLM transformation. Specifically, OPDLM is trained via self-OPD, where the student, an ARLM with bidirectional attention, generates its own trajectories, and the teacher, the original frozen ARLM, distills its knowledge by providing target logits on these trajectories. By training directly in an on-policy manner, OPDLM eliminates the train-inference mismatch in DLMs, while distillation from the original model enhances knowledge retention from the ARLM. Empirical results demonstrate that OPDLM requires 15x to 7,000x fewer training tokens with strong performance across a wide variety of tasks. OPDLM avoids the prohibitive cost of DLM pretraining and positions DLM transformation as a form of ARLM post-training.

1. Introduction … 3. Preliminaries

The paper frames ARLM-to-DLM conversion as data-efficient post-training with OPDLM, addressing knowledge-retention and training–inference mismatches through student-generated diffusion trajectories and frozen-ARLM supervision. It also reviews MDLM masking, reverse decoding, and OPD foundations underlying this approach.

  • Appendix: 0.066B training tokens and 4.2 × 1018 FLOPs are required to train OPDLM-8B, establishing a new pareto frontier for AIME-24.The figure reports these quantities as 15× to 7,000× less than established DLMs obtained from ARLMs.
  • 2. Background and Related Works: MDLMs extend diffusion language modeling to text, offer flexible decoding with potential multi-token updates, and require trillion-scale tokens when pretrained from scratch.ARLM-to-DLM conversion reduces required training data from trillions of tokens to billions, but existing methods largely replace the ARLM loss with the standard DLM loss.
  • 1. Introduction: The method addresses knowledge-retention mismatch from replacing next-token prediction and training–inference mismatch between random masked states and confidence-guided reverse trajectories.These are the two distribution shifts identified for ARLM-to-DLM conversion and standard MDLM training.
  • 1. Introduction: OPDLM converts an ARLM into a DLM by initializing the student from the ARLM, sampling reverse diffusion trajectories, and distilling token-level targets from the frozen original ARLM.This avoids requiring a separately pretrained DLM teacher.
  • 1. Introduction: 15× to 7,000× fewer training tokens are used by OPDLM than established DLM baselines while remaining competitive across model scales and task categories.The paper formulates conversion as post-training, avoiding computationally intensive DLM pretraining.
  • 2. Background and Related Works: On-Policy Distillation trains students on their own sampled trajectories while teachers provide token-level supervision, reducing distribution shift relative to teacher-generated states.Prior self-OPD work shows that a single model can serve as its own teacher in autoregressive settings.
  • 3. Preliminaries: Standard MDLM training samples independently masked states from a forward trajectory, whereas inference follows reverse sampler-induced trajectories whose states depend on model–sampler interaction.This state-distribution difference can make learning less efficient.
  • 3. Preliminaries: OPD compares student and teacher distributions on prefixes induced by the student’s own rollout rather than fixed-dataset or teacher-generated prefixes.For OPDLM, the analogous on-policy states are partially masked diffusion states, with teacher targets conditioned on the corresponding causal prefix.

4. On-Policy Diffusion Language Models

OPDLM converts an ARLM into a DLM by training on the student’s reverse decoding trajectories while distilling the original frozen ARLM’s token-level predictions. This replaces forward-masking training states with inference-aligned states and supports knowledge retention without a separately trained DLM teacher.

  • Motivation: Standard diffusion training uses forward random-masking states, whereas inference uses reverse decoding states, creating the training–inference mismatch that OPDLM directly addresses.The trajectory-level formulation identifies changing the training trajectory distribution as the direct way to reduce this mismatch.
  • On-policy trajectory training: OPDLM replaces forward-masking trajectories with reverse trajectories generated by the student’s sampler, aligning supervised states with the inference process.Training samples a reverse trajectory, selects a realized non-terminal denoising time, and supervises the corresponding state.
  • ARLM-supervised distillation: The frozen original ARLM supplies token-level target distributions, avoiding a separately trained DLM teacher and transferring knowledge through distribution matching.Teacher queries use the terminal sequence from the student trajectory to construct unmasked causal prefixes for each masked position.
  • ARLM-supervised distillation: OPDLM trains the student on its own reverse-trajectory states with the ARLM’s full predictive distributions, encouraging retention of the original model’s knowledge.The method uses uniform loss weighting, w(t) = 1.
  • Rollout curriculum: A rollout curriculum starts with shorter trajectories and gradually exposes the student to longer ones because reverse trajectories can be unstable immediately after conversion.Shorter early rollouts provide contexts primarily comprised of the prompt and avoid weak signals from long, low-quality terminal sequences.

5. Experiments … Appendix

Experiments show that OPDLM converts ARLMs into competitive DLMs with dramatically fewer training tokens, while on-policy data, retained ARLM priors, and controllable decoding improve its practical utility. Ablations, efficiency studies, and task-specific results support OPDLM as an effective post-training approach.

  • 5.1. Setup: OPDLM uses Qwen3 ARLMs from 0.6B to 8B, trains on ∼60K samples across math, code, science, and chat, and evaluates with greedy static decoding at block size=4.The corpus contains 20,222 math, 21,594 code, 10,000 science, and 10,000 chat samples; only prompts are retained for on-policy training.
  • 5.1. Setup: OPDLM matches strong AR-to-diffusion baselines across knowledge, mathematics, and coding benchmarks while using 15× to 7,000× fewer training tokens.Table 1 reports resource use as low as 0.075B tokens.
  • 5.2. OPDLM as a General-Purpose DLM: Across general-purpose benchmarks, OPDLM is competitive while using two to three orders of magnitude fewer tokens than baselines, though SDAR remains stronger on several benchmarks.The comparison is confounded because SDAR uses ∼55B undisclosed training tokens, whereas OPDLM uses ∼66M-76M public-corpus tokens; matched-data and compute results are provided in the ablation.
  • 5.2. OPDLM as a General-Purpose DLM: OPDLM’s advantage grows with task complexity and it retains structural ARLM priors, enabling zero-shot capabilities such as extended thinking and multilingual performance without corresponding training data.It is comparable or superior on GPQA-Diamond, AIME, and LiveCodeBench, while baseline models retain an edge on saturated datasets such as GSM8K.
  • 5.3. Ablation Study: On-Policy vs. Off-Policy Distillation: On-policy generations are the primary performance driver: replacing offline teacher responses with student trajectories consistently improves results at both 4B and 8B scales.The comparison uses matched prompt data for one epoch and keeps ARLM soft targets fixed when moving from Off-Policy to OPDLMoff.
  • 5.3. Ablation Study: On-Policy vs. Off-Policy Distillation: OPDLM removes the training-inference state mismatch by making both masking and target-response generation on-policy, while OPDLMoff isolates the limited effect of the masking trajectory.OPDLM trains on states generated by its own reverse diffusion process; OPDLMoff instead randomly masks the student’s terminal generated sequence.
  • 5.4. Inference Efficiency: Multi-Token Decoding: Decoding throughput is controllable: lowering γ from 1 to 0.8 raises output from 1 token/step to over 2 tokens/step, while increasing block size from 4 to 16 raises throughput from ∼2 to ∼3.5 tokens/step.Both changes incur some accuracy cost; block size sets the upper bound on inference parallelism.
  • 5.5. OPDLM as a Task-Specific DLM: For task-specialized experts, OPDLM can start directly from an AR checkpoint and distill on task-specific data, avoiding a pretrained DLM stage and reward-based intermediate training.In the fair comparison, OPDLM uses the same 8K level 3-5 hard MATH problems as TraDo, but applies on-policy distillation without rewards.

A. Additional Experimental Results · A.1. Results at Smaller Scales · A.2. Effect of Teacher Model Size

At 0.6B and 1.7B scales, OPDLM matches or surpasses comparable baselines on math and reasoning with fewer tokens but underperforms on coding. Teacher-size effects depend on student scale: self-distillation works for larger students, while 0.6B benefits from a larger teacher.

  • A.1. Results at Smaller Scales: OPDLM is evaluated at 0.6B and 1.7B scales against Simple-dLLM, Fast-dLLM-v2, and SDAR.These results are presented in Table 6.
  • A.1. Results at Smaller Scales: At smaller scales, OPDLM matches or surpasses comparable baselines on math and reasoning benchmarks while using significantly fewer tokens.The reported comparisons are against Simple-dLLM, Fast-dLLM-v2, and SDAR.
  • A.1. Results at Smaller Scales: OPDLM underperforms on coding benchmarks where the training data corpus has limited coverage.This coding limitation accompanies the otherwise favorable smaller-scale math and reasoning results.
  • A.2. Effect of Teacher Model Size: The main experiments use self-distillation, pairing a 4B student with a 4B teacher and an 8B student with an 8B teacher.The teacher-size analysis tests whether the optimal teacher depends on student scale.
  • A.2. Effect of Teacher Model Size: For 4B and 8B students, self-distillation matches or outperforms distillation from a stronger Qwen-32B teacher.This finding is reported in Table 7.
  • A.2. Effect of Teacher Model Size: For 0.6B students, self-distillation is insufficient, while a larger Qwen3-4B teacher provides some gains.This contrasts with the results for the larger students in Table 7 and is reported in Table 8.
  • A.2. Effect of Teacher Model Size: The authors hypothesize that a 0.6B teacher lacks sufficiently rich predictive signal on student-generated rollouts.They leave systematic study of teacher–student scaling behavior to future work.

A.3. Additional Results for Multi-token Decoding

Additional experiments vary block size and decoding-confidence threshold across MATH-500, GPQA-Diamond, and MBPP. They confirm a general accuracy-throughput trade-off and show rapid stabilization of training throughput according to block-size-determined parallelism.

  • Experimental setup: Fig. 4 and Fig. 5 sweep block size and decoding-confidence threshold across MATH-500, GPQA-Diamond, and MBPP.The experiments test whether the trade-off extends beyond mathematical reasoning.
  • Accuracy-throughput trade-off: Larger block sizes and lower thresholds produce more tokens per denoising step but reduce accuracy across configurations.This establishes the same accuracy-throughput trade-off across both axes.
  • Training throughput: Throughput rises rapidly during training and then stabilizes, indicating quick convergence to block-size-determined parallelism.Fig. 6 reports the average tokens generated per denoising step over training.
  • Training throughput: Block size 16 maintains the highest throughput, followed by block sizes 8 and 4, while also showing the greatest early-phase variance.The ordering matches inference-time observations.

A.4. Multi-Seed Evaluation on GPQA-Diamond, AIME-24, AIME-25

Because GPQA-Diamond, AIME-24, and AIME-25 are small benchmarks where single-run scores can fluctuate, Table 9 evaluates each model over multiple random seeds and reports mean ± std.

  • Multi-Seed Evaluation: Table 9 reports model accuracy on GPQA-Diamond, AIME-24, and AIME-25 as mean ± std across multiple random seeds.The multi-seed evaluation addresses score fluctuations on these small benchmarks.

A.5. Comparison with Supervised Fine-Tuning

The comparison evaluates supervised fine-tuning (SFT) on ARLM-generated responses against OPDLM under matched data and compute. SFT is an efficient baseline, whereas OPDLM is more robust to complex inference-time samplers and avoids the train–inference divide exposed by dynamic sampling.

  • Method: The ablation trains on one response per prompt using the BD3LM loss for one epoch with block size 4.This extends the off-policy comparison to training directly on ARLM-generated data.
  • Matched comparison: Table 10 compares SFT and OPDLM at the 4B scale using the same prompt corpus and one data epoch.The comparison also matches compute between the two variants.
  • Dynamic sampling: Under dynamic sampling, standard DLM training with random masking shows a much larger average performance drop than OPDLM.Table 11 evaluates static decoding against dynamic sampling with confidence threshold=0.9 and reports per-benchmark changes and AIME standard deviations.
  • Dynamic sampling: SFT fails to address the train-inference divide exposed by dynamic sampling, while OPDLM provides superior robustness across complex inference-time samplers.The results position SFT on ARLM responses as a strong, efficient baseline and OPDLM as the more sampler-robust approach.

B. Experiment Details … B.5. Compute Usage

The experiments use a 62K-sample mixed-domain pretraining corpus, broad benchmark coverage, specified OPDLM hyperparameters, and an explicit FLOP-accounting procedure. Final checkpoints use H200 GPUs, while smaller models can be trained on as few as 2 A6000 GPUs.

  • B.1. Pretraining Datasets Details: 62K samples comprise 22k mathematical-reasoning, 20k coding, 10k STEM, and 10k Chat examples.Mathematics data comes from DAPO and Nemotron-v2-Math; coding data comes from TACO, KodCode-Light-RL, and AceCode.
  • B.2. Evaluation Benchmark Details: Evaluation covers general knowledge, mathematics and reasoning, code generation, and multilingual benchmarks.The listed datasets include MMLU, GSM8K, HumanEval, MMMLU-lite, and additional benchmarks in each category.
  • B.2. Evaluation Benchmark Details: Evaluation uses the official Qwen3 chat template, with task-specific instructions for mathematical reasoning and multiple-choice answers.Mathematics prompts request step-by-step reasoning with a boxed final answer, while multiple-choice prompts request only the choice letter.
  • B.3. Hyperparameters: General-purpose OPDLM hyperparameters are listed in Table 13, while OPDLM-MATH overrides only the parameters specified in Table 14.Non-thinking and thinking variants differ only in rollout length.
  • B.4. FLOPs Calculation: FLOP estimates assign approximately 2N FLOPs per token to forward passes and 6N FLOPs per token to forward-plus-backward passes.The teacher runs forward-only, whereas the student uses forward-plus-backward training; total tokens sum across rollout stages under the curriculum schedule.
  • B.4. FLOPs Calculation: The FLOP calculation conservatively assumes one student forward per token, slightly overestimating OPDLM’s charged FLOPs.In practice, block-diffusion training operates at roughly 1.5 tokens/step rather than 1 token/step.
  • B.5. Compute Usage: Final checkpoints were trained on H200 GPUs, while ablations and intermediate experiments used A6000 and A100 GPUs.All models below 4B parameters can be trained on as few as 2 A6000 GPUs.

C. Limitations, Future Directions and Broader Impacts · Future Directions

The paper identifies data quality, broader reasoning distillation, and fairer comparisons as key areas for improvement, while emphasizing OPDLM’s efficiency and sustainability benefits. Future work also includes scaling reasoning distillation and optimizing decoding and training throughput.

  • C. Limitations, Future Directions and Broader Impacts: OPDLM remains weaker than the strongest baselines on certain tasks despite competitive benchmark performance.The authors identify room for improvement across selected tasks.
  • C. Limitations, Future Directions and Broader Impacts: Publicly unavailable pretraining datasets make controlled comparisons with most baselines difficult.The limitation stems from the lack of access to the datasets used by most baselines.
  • Future Directions: Reasoning distillation should expand beyond the specialized math setting to broader, more general datasets.The paper presents this as future work after exploring thinking-mode distillation in mathematics.
  • Future Directions: A ∼60K-prompt corpus assembled from public sources motivates building higher-quality and more diverse training corpora.The authors expect OPDLM’s data efficiency to make data curation a promising route to substantial gains.
  • Future Directions: Up to 7000× fewer training tokens than from-scratch DLM pre-training lowers the energy footprint of producing competitive generative models.This broader-impact benefit connects efficient ARLM-to-DLM conversion with more sustainable model development.
  • Future Directions: Decoding performance is studied across block sizes 4, 8, and 16 and thresholds γ ∈{0.95, 0.90, 0.85, 0.80} on MATH-500, GPQA-Diamond, and MBPP.Figures 4 and 5 examine the effects of block size and decoding threshold across these benchmarks.
  • Future Directions: Larger block sizes generate more tokens per denoising step, with block size 16 achieving the highest average token throughput during training.Figure 6 reports training curves for average tokens generated per denoising step.
Loading 2606.06712v1…