TAD: Temporal-Aware Trajectory Self-Distillation for Fast and Accurate Diffusion LLM

TL;DR

TAD introduces a temporal-aware self-distillation framework for diffusion LLMs, enhancing the accuracy-parallelism trade-off by leveraging decoding trajectory information during training.

Contribution

It proposes a novel trajectory self-distillation method that partitions tokens based on decoding steps, improving both speed and accuracy in diffusion LLMs.

Findings

Increases LLaDA accuracy from 46.2% to 51.6%.

Raises average AUP from 46.2 to 257.1 with the Speed model.

Consistently improves the accuracy-parallelism trade-off.

Abstract

Diffusion large language models (dLLMs) offer a promising paradigm for parallel text generation, but in practice they face an accuracy-parallelism trade-off, where increasing tokens per forward (TPF) often degrades generation quality. Existing acceleration methods often gain speed at the cost of accuracy. To address this limitation, we propose TAD, a Temporal-Aware trajectory self-Distillation framework. During data construction, we condition a teacher model on both the prompt and the ground-truth response to generate decoding trajectories, recording the intermediate masked states throughout the process. Based on how many decoding steps remain before each masked token is revealed, we partition masked positions into near and distant subsets. For near tokens, we train the student with a hard cross-entropy loss using the teacher trajectory tokens as labels, encouraging confident predictions for tokens that are about to be decoded. For distant tokens, we apply a soft KL divergence loss between the teacher and student token distributions, providing softer supervision and preserving future planning knowledge. This temporal-aware partition naturally gives rise to two deployment configurations: a Quality model that prioritizes accuracy and a Speed model that favors more aggressive acceleration. Experiments show that TAD consistently improves the accuracy-parallelism trade-off. On LLaDA, it raises average accuracy from 46.2\% to 51.6\% with the Quality model and average AUP from 46.2 to 257.1 with the Speed model. Our code is available at: https://github.com/BHmingyang/TAD