TAD：时间感知轨迹自蒸馏，用于快速准确的扩散 LLM

摘要

Diffusion 大语言模型 (dLLM) 为并行文本生成提供了一种有前途的范例，但在实践中，它们面临着准确性与并行性的权衡，其中增加每个转发的词元 (TPF) 通常会降低生成质量。现有的加速方法常常以牺牲准确性为代价来获得速度。为了解决这个限制，我们提出了 TAD，一种时间感知轨迹自蒸馏框架。在数据构建过程中，我们根据提示和真值响应来调节教师模型，以生成解码轨迹，记录整个过程中的中间屏蔽状态。根据每个屏蔽标记被揭示之前剩余的解码步骤数，我们将屏蔽位置划分为近子集和远子集。对于近标记，我们使用教师轨迹标记作为标签来训练具有硬交叉熵损失的学生，鼓励对即将解码的标记进行自信的预测。对于远程词元，我们在教师和学生词元分布之间应用软 KL 散度损失，提供更软的监督并保留未来的规划知识。这种时间感知分区自然会产生两种部署配置：优先考虑准确性的质量模型和支持更积极加速的速度模型。实验表明，TAD 持续改善了准确性与并行性的权衡。在 LLaDA 上，质量模型的平均准确率从 46.2% 提高到 51.6%，速度模型的平均 AUP 从 46.2 提高到 257.1。我们的代码位于：https://github.com/BHmingyang/TAD