Few-Step Diffusion Language Models via Trajectory Self-Distillation

TL;DR

This paper introduces a self-distillation framework for diffusion large language models that significantly improves few-step decoding quality, enabling faster text generation without substantial performance loss.

Contribution

It proposes a trajectory-level supervision method combined with Direct Discriminative Optimization to enhance few-step decoding in diffusion language models, reducing quality degradation.

Findings

Substantially narrows the gap between few-step and full-step decoding.

Improves performance on reasoning and code-generation benchmarks.

The source code is publicly available at https://github.com/Tyrion58/T3D.

Abstract

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.