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

2026-05-10

AI Summary

This work addresses the degradation in generation quality often observed when accelerating parallel decoding in diffusion-based large language models. To mitigate this trade-off, the authors propose the Temporal-Aware Distillation (TAD) framework, which leverages a teacher model to generate decoding trajectories during data construction. Masked positions are categorized as “near” or “far” based on their remaining decoding steps, and are supervised with hard cross-entropy loss and soft KL divergence loss, respectively, enabling temporally aware self-distillation. This approach introduces a novel temporal-aware mechanism that provides stage-specific supervision, supporting both accuracy-prioritized and speed-prioritized deployment modes. Experiments on LLaDA demonstrate that the accuracy-prioritized variant improves average accuracy from 46.2% to 51.6%, while the speed-prioritized model achieves a substantial increase in average AUP to 257.1, effectively alleviating the accuracy–speed trade-off in parallel generation.

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

Problem

• diffusion LLM
• accuracy-parallelism trade-off
• parallel text generation
• token generation quality

Innovation

• temporal-aware distillation
• diffusion LLM
• trajectory self-distillation
• parallel text generation
• accuracy-parallelism trade-off