通过时间感知自蒸馏,让扩散语言模型在提速的同时保持高精度。
TAD: Temporal-Aware Trajectory Self-Distillation for Fast and Accurate Diffusion LLM

- 根据解码进度将掩码位置分为近/远两类,分别用硬标签和软监督训练。
- 质量模型准确率提升至51.6%,速度模型平均生成效率提高4.5倍。
- 适用于需要平衡生成速度与质量的文本生成场景。
扩散大语言模型(dLLMs)为并行文本生成提供了有前景的范式,但实践中面临准确率与并行性的权衡:增加每前向传播的标记数(TPF)常导致生成质量下降。现有加速方法往往以牺牲准确率为代价。为此,我们提出TAD——一种时间感知轨迹自蒸馏框架。在数据构建阶段,将教师模型同时基于提示和真实响应进行条件化,生成解码轨迹,并记录全过程中的中间掩码状态。根据每个掩码标记即将被揭示的解码步数,将掩码位置划分为近距和远距子集。对近距标记,使用教师轨迹标记作为硬标签,施加交叉熵损失,鼓励学生模型对即将解码的标记做出自信预测;对远距标记,采用教师与学生标记分布间的软KL散度损失,提供更柔和的监督并保留未来规划知识。该时间感知划分自然催生两种部署配置:侧重准确性的质量模型和追求更高加速的速模。实验表明,TAD持续改善了准确率-并行性权衡。在LLaDA上,质量模型使平均准确率从46.2%提升至51.6%,速度模型使平均AUP从46.2提升至257.1。代码已开源:https://github.com/BHmingyang/TAD
原文摘要 · Abstract (English)
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
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。