Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

T3D: Few-Step Diffusion Language Models via Trajectory Self-Distillation with Direct Discriminative Optimization

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

TL;DR

This work tackles the latency bottleneck of diffusion-based language models by introducing T3D, a trajectory self-distillation framework that trains a few-step student from teacher rollout trajectories using a mode-seeking Direct Discriminative Optimization objective. By combining trajectory-level supervision with a discriminative, likelihood-ratio loss and an on-policy training regime, T3D reduces factorization error in the reverse diffusion and stabilizes learning under tight step budgets. Theoretical analysis shows trajectory distillation lowers conditional dependence and TC, while experiments across reasoning and code-generation benchmarks demonstrate improved few-step accuracy and preserved full-step performance, even under dynamic decoding. Overall, T3D narrows the gap to full-step diffusion, enabling efficient, practical few-step DLLMs with broad applicability.

Abstract

Diffusion large language models (DLLMs) have the potential to enable fast text generation by decoding multiple tokens in parallel. However, in practice, their inference efficiency is constrained by the need for many refinement steps, while aggressively reducing the number of steps leads to a substantial degradation in generation quality. To alleviate this, we propose a trajectory self-distillation framework that improves few-step decoding by distilling the model's own generative trajectories. We incorporate Direct Discriminative Optimization (DDO), a reverse-KL objective that promotes mode-seeking distillation and encourages the student to concentrate on high-probability teacher modes. Across benchmarks, our approach consistently outperforms strong few-step baselines and standard training under tight step budgets. Although full-step decoding remains superior, we substantially narrow the gap, establishing a strong foundation towards practical few-step DLLMs. The source code is available at https://github.com/Tyrion58/T3D.

T3D: Few-Step Diffusion Language Models via Trajectory Self-Distillation with Direct Discriminative Optimization

TL;DR

This work tackles the latency bottleneck of diffusion-based language models by introducing T3D, a trajectory self-distillation framework that trains a few-step student from teacher rollout trajectories using a mode-seeking Direct Discriminative Optimization objective. By combining trajectory-level supervision with a discriminative, likelihood-ratio loss and an on-policy training regime, T3D reduces factorization error in the reverse diffusion and stabilizes learning under tight step budgets. Theoretical analysis shows trajectory distillation lowers conditional dependence and TC, while experiments across reasoning and code-generation benchmarks demonstrate improved few-step accuracy and preserved full-step performance, even under dynamic decoding. Overall, T3D narrows the gap to full-step diffusion, enabling efficient, practical few-step DLLMs with broad applicability.

Abstract

Diffusion large language models (DLLMs) have the potential to enable fast text generation by decoding multiple tokens in parallel. However, in practice, their inference efficiency is constrained by the need for many refinement steps, while aggressively reducing the number of steps leads to a substantial degradation in generation quality. To alleviate this, we propose a trajectory self-distillation framework that improves few-step decoding by distilling the model's own generative trajectories. We incorporate Direct Discriminative Optimization (DDO), a reverse-KL objective that promotes mode-seeking distillation and encourages the student to concentrate on high-probability teacher modes. Across benchmarks, our approach consistently outperforms strong few-step baselines and standard training under tight step budgets. Although full-step decoding remains superior, we substantially narrow the gap, establishing a strong foundation towards practical few-step DLLMs. The source code is available at https://github.com/Tyrion58/T3D.
Paper Structure (37 sections, 6 theorems, 49 equations, 1 figure, 7 tables, 1 algorithm)

This paper contains 37 sections, 6 theorems, 49 equations, 1 figure, 7 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Masked Diffusion Models
  4. Direct Discriminative Optimization (DDO)
  5. Methods
  6. Trajectory Self-Distillation
  7. Trajectory Distillation via DDO
  8. Path Consistency Regularization
  9. T3D: Trajectory Self-Distillation via DDO
  10. Theoretical Analysis
  11. Self-Distillation Setting
  12. Optimality of Trajectory Self-Distillation
  13. Trajectory-Distillation as Discrete Rectified Flow.
  14. Rectified Discrete Flow (ReDi).
  15. Trajectory Distillation induces lower Conditional Total Correlation.
  16. ...and 22 more sections

Key Result

Proposition 4.3

For a teacher model $p_\phi$ and student model $p_\theta$, we define the on-policy risk for distillation: $\forall t \in [T]$, Given the target joint distribution $J_\phi$ and $J_\phi^\mathrm{Tra}$ defined in Eqn. eq: def-joint-dis and Definition def: traj-joint, as well as joint distribution of student model $J_\theta$ in Eqn. eq: def-joint-dis, let Under the Assumption asp:margin, there is

Figures (1)

  • Figure 1: Overview of T3D for enabling few-step diffusion decoding.(a) Teacher full trajectory. A pretrained diffusion language model generates sequences through a full diffusion trajectory from the fully masked state ${\bm{x}}_T$ to the clean sequence ${\bm{x}}_0$ using many refinement steps, producing intermediate states that reflect the inference-time decoding distribution. (b) Trajectory-level objective learning. T3D trains a few-step student model via trajectory self-distillation on teacher rollout trajectories. We incorporate Direct Discriminative Optimization (DDO) to perform mode-seeking trajectory matching between the student and a reference model under an on-policy discriminative objective. (c) Few-step decoding. After training, the student can decode using significantly fewer diffusion steps (e.g., 1--2 steps per block) while preserving generation quality, substantially narrowing the gap with full-step decoding and enabling efficient parallel token generation.

Theorems & Definitions (11)

  • Definition 4.1: Trajectory Joint Distribution
  • Proposition 4.3: Optimality of $J_\phi^\mathrm{Tra}$ for on-policy posterior matching
  • Definition 4.4: Conditional Total Correlation (TC)
  • Theorem 4.5: Trajectory distribution induces lower conditional dependence
  • Lemma 2.2: KL decomposition for joint distributions
  • Proposition 2.3: Optimality of $J_\phi^\mathrm{Tra}$ for on-policy posterior matching
  • proof
  • Definition 2.4: Conditional Total Correlation (TC)
  • Lemma 2.7: Pythagorean Inequality for KL Divergence wolfer2024geometric
  • Theorem 2.8: Trajectory distribution induces lower conditional dependence
  • ...and 1 more