Few-Step Diffusion Language Models via Trajectory Self-Distillation
About
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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Mathematical Reasoning | MATH500 (test) | Accuracy61.6 | 895 | |
| Instruction Following | IFEval | -- | 836 | |
| Code Generation | HumanEval (test) | -- | 612 | |
| Code Generation | MBPP (test) | -- | 405 | |
| Radiology Report Generation | MIMIC-CXR (test) | ROUGE-L52.38 | 209 | |
| Radiology Report Generation | CheXpert Plus (test) | -- | 88 | |
| Mathematical Reasoning | GSM8K (test) | Accuracy0.8385 | 48 | |
| Code Generation | HumanEval | TPS222.7 | 41 | |
| Mathematical Reasoning | MATH 500 | Overall Score40.7 | 25 | |
| Chest X-ray Report Generation | ReXGradient (test) | ROUGE-L64.36 | 16 |