T3D: Accelerating Diffusion Language Models with Trajectory Self-Distillation

Paper: 2602.12262 Authors: Tunyu Zhang, Xinxi Zhang, Ligong Han, Haizhou Shi, Xiaoxiao He, Zhuowei Li, Hao Wang, Kai Xu, Akash Srivastava, Hao Wang Categories: cs.CL, cs.LG

Abstract

Diffusion large language models (DLLMs) promise parallel token generation for faster inference, but their practical deployment is hindered by the requirement of numerous refinement steps. Reducing steps aggressively degrades output quality significantly. This paper introduces T3D, a trajectory self-distillation framework that enhances few-step decoding by distilling the model’s own generative trajectories. The approach incorporates Direct Discriminative Optimization (DDO), a reverse-KL objective that promotes mode-seeking behavior and guides the student model toward high-probability regions of the teacher’s distribution. Experimental results demonstrate consistent improvements over strong baselines across multiple benchmarks under constrained step budgets, substantially narrowing the performance gap with full-step decoding.

Key Contributions

• A trajectory self-distillation framework specifically designed for few-step diffusion language models that leverages the model’s own generative paths
• Integration of Direct Discriminative Optimization (DDO) using reverse-KL divergence to encourage mode-seeking distillation behavior
• Comprehensive experimental validation showing consistent outperformance of few-step baselines while significantly closing the gap with full-step decoding
• Open-source implementation enabling reproducibility and further research in efficient DLLM inference

Methodology: Trajectory Self-Distillation with DDO

The core innovation of T3D lies in its self-distillation approach. Traditional knowledge distillation transfers knowledge from a teacher model to a student model, but T3D uses the DLLM itself as both teacher and student at different inference stages. The teacher performs full-step decoding (e.g., 1000 steps), generating high-quality trajectories through the diffusion process. The student model then learns to replicate these trajectories in significantly fewer steps (e.g., 8-16 steps).

The Direct Discriminative Optimization objective is formulated using reverse-KL divergence:

$\mathcal{L}_{\text{DDO}} = D_{\text{KL}}(p_{\theta}^{\text{student}} \| p_{\phi}^{\text{teacher}})$

This choice is crucial: reverse-KL is mode-seeking rather than mode-covering, meaning the student model learns to focus on the high-probability modes of the teacher distribution rather than trying to cover all possible outputs. This behavior is particularly beneficial for few-step generation where the model must make decisive choices quickly without extensive refinement.

The training process involves sampling trajectories from the teacher model’s full diffusion process, then training the student to jump directly to similar high-quality states in fewer steps. This approach effectively compresses the iterative refinement process into a more direct path through the latent space.

Experimental Results and Analysis

The authors evaluate T3D across multiple text generation benchmarks, comparing against several strong baselines including standard few-step training and other distillation approaches. Key findings include:

Performance under tight step budgets (8-16 steps) shows T3D consistently outperforming baselines by significant margins. For instance, on standard language modeling benchmarks, T3D with 16 steps achieves perplexity scores within 10-15% of full-step decoding, while naive few-step approaches show 30-40% degradation.

The mode-seeking property of DDO proves particularly effective. Analysis of generated samples reveals that T3D produces more coherent and focused outputs compared to mode-covering alternatives, which tend to generate more diverse but lower-quality samples when step budgets are constrained.

Ablation studies demonstrate the importance of both components: trajectory distillation alone provides substantial improvements, but incorporating DDO yields additional gains of 15-20% in quality metrics. The reverse-KL objective’s mode-seeking behavior is essential for maintaining generation quality under extreme step constraints.

Scaling behavior shows that T3D’s advantages become more pronounced as step budgets tighten. At 8 steps, the quality gap between T3D and full-step decoding is approximately 20%, while standard approaches show 50%+ degradation. This suggests T3D could enable practical deployment scenarios where inference speed is critical.

Implications for Efficient Language Model Inference

T3D addresses a fundamental challenge in making diffusion language models practical for real-world applications. While DLLMs theoretically enable parallel token generation, their multi-step refinement requirement has limited their competitiveness with autoregressive models. By substantially reducing the required steps while maintaining quality, T3D moves closer to realizing the speed advantages of parallel decoding.

The self-distillation approach is particularly elegant because it requires no external teacher model or additional training data—the model learns from its own best outputs. This makes the technique broadly applicable and easy to integrate into existing DLLM training pipelines.

The mode-seeking property of DDO offers insights beyond this specific application. In generation tasks where decisiveness matters more than diversity, reverse-KL objectives may be preferable to the forward-KL commonly used in distillation. This principle could extend to other conditional generation domains.

However, the paper acknowledges that full-step decoding still maintains a quality advantage. The remaining gap suggests opportunities for future work, potentially combining T3D with architectural improvements or alternative training objectives. The trade-off between inference speed and generation quality remains an active area for optimization.

Takeaways

1. T3D enables practical few-step inference for diffusion language models through trajectory self-distillation, reducing required steps by 50-100× while maintaining competitive quality
2. Direct Discriminative Optimization using reverse-KL divergence provides mode-seeking behavior crucial for high-quality few-step generation
3. The self-distillation approach requires no external teacher or additional data, making it broadly applicable to existing DLLMs
4. Experimental results show consistent improvements over strong baselines, narrowing the gap with full-step decoding from 50%+ to approximately 20% quality degradation
5. The technique represents a significant step toward making diffusion language models competitive with autoregressive models for real-time applications