PaperlayerView exact revision on arXiv

Source-linked AI summary

Few-Step Diffusion Language Models via Trajectory Self-Distillation

Tunyu Zhang, Xinxi Zhang, Ligong Han, Haizhou Shi, Xiaoxiao He, Zhuowei Li, Hao Wang, Kai Xu, Akash Srivastava, Chengzhi Mao, Hao Wang, Vladimir Pavlovic, Dimitris N. Metaxas

arXiv:2602.12262v4cs.CLcs.LG

TL;DR

Few-step DLLM decoding suffers from token-factorization error, making aggressive step reduction degrade generation quality. T3D distills a full-step teacher’s trajectory into a few-step student and adds DDO for sharper mode-seeking alignment. Across reasoning and code-generation benchmarks, it consistently improves prior few-step methods and narrows the gap to full-step decoding, with teacher quality and offline rollout cost as key limitations.

  • Problem

    Reducing DLLM decoding steps degrades output quality because token-factorization error becomes more severe during aggressive jumps.

  • Method

    T3D trains a few-step student to match a full-step teacher’s intermediate generative trajectory and uses DDO to encourage alignment with teacher-preferred modes.

  • Results

    Across reasoning and code-generation benchmarks, T3D consistently outperforms prior few-step DLLM methods and substantially narrows the gap to full-step diffusion decoding.

  • Takeaways & Limitations

    Trajectory self-distillation provides a practical route toward efficient few-step diffusion language modeling across decoding budgets and model families.

  • Takeaways & Limitations

    Student performance is bounded by teacher quality, while collecting full-step teacher trajectories adds offline cost that scales with dataset size and decoding budget.

Abstract

from arXiv · show

Diffusion large language models (DLLMs) have emerged as powerful generative models with the promise of fast text generation through parallel decoding. However, realizing this potential in practice remains challenging: reducing the number of decoding steps, typically causes a substantial degradation in output quality due to token factorization error. To alleviate this, we propose a self-distillation framework that trains a few-step student to match the generative trajectory of a full-step teacher. We theoretically and empirically show that trajectory-level supervision mitigates this factorization error, thereby enabling effective few-step decoding. We further incorporate Direct Discriminative Optimization (DDO), a reverse-KL objective that encourages mode-seeking toward the teacher's modes, yielding stronger performance on challenging reasoning tasks. Across reasoning and code-generation benchmarks, our method substantially narrows the gap between few-step and full-step decoding. The source code is available at https://github.com/Tyrion58/T3D.

1 Introduction

Few-step decoding in diffusion language models is limited by token-factorization error, which worsens as decoding becomes more aggressive. T3D addresses this by matching full-step teacher trajectories and adding DDO-based trajectory alignment.

  • Motivation: Reducing diffusion decoding steps makes simultaneous multi-token generation less accurate, limiting the efficiency gains of DLLMs.This is especially problematic for real-time and compute-constrained applications.
  • Factorization Error: Few-step decoding is fundamentally constrained by mean-field token factorization, whose approximation error increases with larger decoding jumps.Larger gaps between s and t amplify the error and can produce degenerate generations.
  • Trajectory Self-Distillation: Trajectory Self-Distillation trains a few-step student to match the full generative trajectory of its teacher rather than only the endpoint.Intermediate trajectory supervision exposes richer information about the teacher’s prediction structure.
  • Trajectory Self-Distillation: Trajectory-level supervision reduces factorization error across intermediate reverse transitions, avoiding the uninformative endpoint-supervision failure mode in masked diffusion models.The masked prior is deterministic, so endpoint-based rectified-flow-style supervision does not reduce the relevant error.
  • Trajectory Alignment: DDO encourages the student to focus on teacher high-probability modes, while path consistency emphasizes early decoded tokens.These choices target sharper predictions and error propagation under tight decoding budgets.
  • Results: Across reasoning and code-generation benchmarks, T3D improves prior few-step DLLM methods, especially under aggressive budgets, while preserving full-step and dynamic-decoding performance.The evaluation covers SDAR and LLaDA across multiple decoding budgets and model families.

2 Related Work

Prior diffusion-model acceleration work reduces sampling cost through consistency, flow-map, system, or decoding improvements. This work extends the trajectory-distillation perspective to discrete diffusion language modeling and targets the model’s internal prediction structure.

  • Few-Step Diffusion: Consistency and flow-map methods accelerate diffusion generation by enforcing cross-time consistency or directly modeling state-to-state displacements.Distillation variants often use teacher trajectories, including intermediate states or rectified trajectories.
  • Few-Step Diffusion: The paper brings trajectory-based distillation from continuous diffusion methods to discrete diffusion language modeling.This positions trajectory supervision as the bridge between prior diffusion acceleration strategies and MDLMs.
  • Efficient Inference for DLLMs: DLLM efficiency research includes system and decoding advances, while another line directly reduces the number of sampling steps.Examples include KV caching, dynamic decoding, block-structured generation, and step-reduction methods.

3 Background

MDLMs generate discrete token sequences by independently masking tokens in a forward process and learning a token-factorized reverse denoising process. Their training objective reduces to masked-token cross-entropy, while DDO supplies a likelihood-based discriminative objective.

  • Masked Diffusion Language Models: MDLMs are diffusion-based generative models for discrete text sequences, with token sequences drawn from a vocabulary augmented by a mask token.The data distribution produces length-L sequences of discrete tokens.
  • Masked Diffusion Language Models: The forward diffusion process independently masks tokens over continuous time according to a decreasing noise schedule.With α_t = 1 − t, each token is preserved with probability α_t and masked otherwise.
  • Masked Diffusion Language Models: The reverse denoising process recovers a cleaner sequence from a noisier one and is approximated by a neural network that factorizes across tokens.The reverse transition maps x_t to x_s for s < t.
  • Masked Diffusion Language Models: Maximizing the MDLM evidence lower bound reduces to a masked-token cross-entropy objective.This provides the basic training formulation for the masked diffusion model.
  • Direct Discriminative Optimization: DDO is a GAN-inspired objective that distinguishes data samples from reference-model samples using likelihood ratios rather than an additional discriminator network.The reference model supplies the fake samples, and the discriminator is implicitly parameterized through model likelihoods.
  • Approximation Error: Conditional total correlation rises with larger decoding intervals, revealing higher factorization error during more aggressive few-step jumps.T3D lowers this error surface by matching the teacher trajectory across intervals.

4 Methods

T3D addresses few-step MDLM degradation by training a student on teacher trajectories rather than only endpoints, reducing factorization error under compressed decoding. It combines trajectory matching with DDO and early-token path weighting to sharpen predictions and preserve important dependencies.

  • Factorization Error: Conditional Total Correlation rises with larger decoding intervals, identifying token-factorization error as the central bottleneck in few-step MDLMs.Larger jumps require modeling stronger cross-token dependencies, while tokenwise factorization becomes increasingly inaccurate.
  • Trajectory Self-Distillation: Trajectory self-distillation trains a few-step student directly on teacher rollout trajectories and exposes it to prediction structure throughout the reverse process.This provides richer supervision than endpoint-only distillation and targets the source of few-step failure.
  • Trajectory Self-Distillation: Trajectory-level supervision reduces approximation error, whereas endpoint-only distillation cannot reduce Conditional Total Correlation in MDLMs with a deterministic prior.The student operates on intermediate states, yielding a substantially lower TC surface across decoding intervals.
  • DDO: DDO replaces forward KL with reverse-KL-like mode-seeking that focuses the student on teacher high-probability trajectories and produces sharper refinement after broad initial exploration.On MATH500 with SDAR-4B-Chat, DDO maintains higher entropy at the fully masked stage but substantially lower entropy afterward.
  • Path-Consistency Regularization: Early-token path-consistency weighting consistently outperforms uniform and late-token weighting on MATH500 and GSM8K under tight step budgets.The method assigns greater weight to early-decoded tokens because their errors are more likely to propagate through the remaining trajectory.
  • Full Method: T3D combines teacher-generated trajectories, DDO, and path-consistency regularization to support efficient few-step decoding while remaining among the strongest methods in SDAR comparisons.Few-step performance is evaluated using tokens-per-step across SDAR-1.7B-Chat and SDAR-4B-Chat.

5 Experiments

Experiments show that T3D performs strongly across few-step, full-step, and dynamic decoding settings. It improves few-step accuracy, preserves full-diffusion performance, and gains efficiency under adaptive decoding.

  • Few-step decoding: T3D achieves the best average accuracy on LLaDA at both TokPS = 4 and TokPS = 8.
  • Few-step decoding: T3D is consistently among the strongest self-distillation methods across SDAR and LLaDA few-step evaluations.
  • Full decoding: T3D achieves performance nearly identical to the original pretrained model under full decoding across SDAR-1.7B-Chat and SDAR-4B-Chat benchmarks.
  • Full decoding: T3D preserves strong full-step performance, whereas ReDi and dParallel exhibit substantial degradation after reverting to full diffusion decoding.
  • Dynamic decoding: +10.8% absolute accuracy and 1.45× higher throughput are achieved by T3D over the original model under dynamic GSM8K decoding.T3D also decodes more tokens per step under the same confidence threshold, suggesting more confident predictions for adaptive decoding.
  • Ablation study: Trajectory distillation provides the main ablation improvement, while DDO and path consistency add further gains under aggressive few-step decoding.The path-consistency loss emphasizes early decoded tokens and helps reduce error propagation.

6 Conclusion

The conclusion identifies factorization error as the bottleneck for few-step MDLM decoding and presents trajectory self-distillation as a richer supervision strategy. T3D outperforms prior few-step methods across reasoning and code-generation benchmarks while retaining teacher- and rollout-related limitations.

  • Conclusion: Few-step MDLM decoding is bottlenecked by factorization error, which trajectory supervision targets using the teacher’s full generative trajectory.
  • Conclusion: T3D distills intermediate decoding intervals where the reverse process remains informative, avoiding a failure mode of endpoint-based self-distillation.
  • Conclusion: T3D consistently outperforms prior few-step DLLM methods across reasoning and code-generation benchmarks, substantially narrowing the gap to full-step diffusion decoding.
  • Limitations: Trajectory collection requires full-step teacher rollouts, so its offline cost scales with dataset size and decoding budget.
  • Limitations: Student performance is ultimately bounded by teacher quality because T3D relies on self-distillation.

B Proof of Theoretical Analysis

The theoretical analysis studies how trajectory self-distillation affects factorization properties in few-step MDLM decoding. It establishes a lower conditional total-correlation result for trajectory distillation and explains why endpoint-only distillation cannot reduce the corresponding quantity.

  • Proof setup: The analysis focuses on trajectory self-distillation and its effect on factorization properties of the reverse diffusion process under few-step decoding.
  • Theoretical result: The proof relies on an optimal MDLM student assumption, a log-convex-hull assumption, and a Pythagorean inequality for KL divergence.
  • Theoretical result: Theorem B.4 states that trajectory distillation induces lower conditional total correlation under its stated assumptions.
  • Endpoint-only distillation: Endpoint-only distillation does not reduce conditional total correlation because the MDLM prior p(xT) = δm is deterministic.
  • Endpoint-only distillation: With a deterministic prior, p(x0 | xT) = q(x0), making the conditional quantity a fixed constant of the data distribution.

C.5 Other Implementation Details

Implementation details cover baselines, data and trajectory construction, training cost, and additional dynamic-decoding experiments. T3D remains effective under adaptive decoding despite being trained with static budgets.

  • Experimental setup: T3D is compared with ReDi, dParallel, SFT, and CDLM for LLaDA under the paper’s few-step evaluation settings.
  • Training data: Self-distillation training uses model-generated MATH and PrimeIntellect responses, while SFT uses data derived from Bespoke-Stratos-17k.
  • Trajectory construction: Teacher trajectories are collected with low-confidence remasking and static decoding, recording token decoding order to reconstruct intermediate states.
  • Training cost: For SDAR-4B-Chat, trajectory collection takes approximately 1.5 hours and T3D training approximately 8 hours on 8× NVIDIA A100-40GB GPUs.
  • Dynamic decoding: T3D improves both accuracy and throughput on MATH500 and GSM8K under dynamic decoding while maintaining competitive code-generation accuracy.
  • Dynamic decoding: These dynamic-decoding results support compatibility with adaptive strategies even though T3D is trained under static decoding budgets.

D.2 Useful Exploration under Reverse-KL Training

T3D’s reverse-KL training preserves useful exploration rather than causing uniform diversity collapse. Its benefits extend from structured reasoning and coding to broader language tasks, with stable multi-seed performance.

  • T3D shows higher entropy early and lower entropy later, combining broader exploration with sharper refinement during decoding.This stage-wise exploration–exploitation pattern is described as desirable for reasoning.
  • T3D outperforms Forward-KL at every pass@k value, while its gap to the teacher narrows as k increases.The results indicate preserved exploration ability and benefits from larger pass@k budgets.
  • 46.27 average accuracy on MATH-500 with 0.31 standard deviation demonstrates strong, low-variance performance across three seeds.On MBPP, T3D reaches 21.87 average accuracy and outperforms the original model and prior few-step DLLM baselines.
  • T3D outperforms dParallel and ReDi on WinoGrande, extending its few-step benefits beyond math and coding benchmarks.WinoGrande tests commonsense language understanding and broader language capability.

E.1 The Effectiveness of λ in Training Objective

Moderate path-consistency regularization is most effective across the tested configurations, while direct few-step distillation can damage full-step behavior. The full T3D objective restores and slightly exceeds the original full-step accuracy.

  • λ = 0.2 achieves the strongest or near-best performance in most MATH500 configurations, especially under aggressive few-step decoding.The ablation varies Tokens Per Step, block size, and decoding steps.
  • λ = 0.05 is often insufficient to stabilize training, whereas λ = 0.5 degrades performance in several settings.The authors therefore fix λ = 0.2 for the main experiments.
  • Few-step distillation alone substantially degrades full-step diffusion behavior when the model is reverted to the original decoding process.Both TD and DDO alone perform poorly, indicating damage to fine-grained denoising capability.
  • 69.00 accuracy from the full T3D objective slightly surpasses the original model under full-step decoding.Adding path consistency raises TD accuracy from 22.00 to 58.00, showing improved retention of intermediate denoising behavior.

E.3 Ablation Study on Few-Step Generation

Trajectory-level distillation provides the main few-step generation gain, especially under the most aggressive decoding budget. DDO further improves trajectory matching, while path consistency stabilizes compressed decoding.

  • Trajectory-level distillation improves over the original model under both tested budgets, with the largest gain at 2 decoding steps per block.The 2-step setting is the more aggressive regime for block size 8 on MATH500.
  • 16.80 to 38.80 accuracy under 2 decoding steps per block shows the gain from trajectory-level distillation.The result supports reducing factorization error caused by large denoising jumps.
  • 45.00 accuracy under the aggressive setting makes the full T3D objective the best-performing ablation.DDO replaces forward-KL trajectory matching with a mode-seeking objective, and its benefit becomes clearer as the decoding budget tightens.
  • Trajectory supervision supplies the main gain, DDO improves trajectory-matching quality, and path consistency further stabilizes compressed decoding.
Loading 2602.12262v4…