基于时空并行解码与置信度外推的高效扩散大语言模型

背景

基于扩散的大语言模型（Diffusion-based Large Language Models, dLLMs）通过迭代去噪过程支持文本的并行生成，这一特性使其在理论上具备超越自回归模型（Autoregressive Models）的并行化潜力。然而，尽管生成方式不同，dLLMs 在实际推理过程中仍然面临显著的延迟问题。

主要瓶颈在于：在去噪的许多步骤中，模型花费了大量计算资源对冗余信息进行微调，或者对最终值已经确定的 token 进行重复的重新掩码（remasking）处理。这种低效的资源分配导致了推理速度的下降。

现有的加速方法主要依赖于两个局限性的策略：

1. 基于步骤的局部置信度启发式规则：这类方法通常对提示词（prompt）和任务变化非常敏感，泛化能力较差。
2. 固定调度策略：忽略了序列内部强烈的位置效应（positional effects），无法动态适应不同 token 的收敛状态。

因此，如何精准判断 token 何时收敛、何时可以安全固定，从而减少不必要的去噪迭代，成为提升 dLLM 推理效率的关键挑战。

核心内容

本文提出了一种将扩散解码视为动态控制问题的新视角，并指出 token 级别的去噪轨迹（denoising trajectories）提供了实现可靠控制的关键信号。基于此，作者提出了一个名为 Trace-aware Decoding Framework（轨迹感知解码框架）的新方法，该框架包含两个核心组件：时空并行解码（TSPD）和置信度外推（CE）。

1. 时空并行解码 (Temporal-Spatial Parallel Decoding, TSPD)

TSPD 的核心在于引入一个轻量级的“时空控制器”（temporal-spatial controller）。该控制器通过综合以下特征来决定一个 token 是否已经收敛，从而可以安全地固定其值，不再参与后续的去噪步骤：

• 每 token 的轨迹特征：
  • 置信度 (Confidence)：模型对当前 token 预测值的确定程度。
  • 熵 (Entropy)：预测分布的不确定性度量。
  • 动量 (Momentum)：token 值在去噪过程中的变化趋势和稳定性。
• Token 位置 (Token Position)：考虑到序列中不同位置的信息依赖性和收敛特性差异。

通过整合时间维度（轨迹动态）和空间维度（位置信息），TSPD 能够更准确地识别出哪些 token 已经稳定，进而跳过对其的冗余计算。

2. 置信度外推 (Confidence Extrapolation, CE)

为了解决轨迹震荡或置信度不足导致决策滞后或错误的问题，作者引入了 Confidence Extrapolation (CE)。这是一个无需训练（training-free）的状态空间模块，其功能包括：

• 预测未来 Logit 趋势：CE 利用不确定性估计，前瞻性地预测 token 在后续去噪步骤中的 logits 变化趋势。
• 支持主动决策：
  • 安全的前瞻（Safe Look-ahead）：在轨迹尚未完全稳定时，通过预测判断其最终收敛方向，允许提前固定 token。
  • 针对性稳定（Targeted Stabilization）：当检测到轨迹处于震荡状态或模型置信度较低时，CE 可以提供额外的信号以辅助稳定决策，避免过早固定导致的质量下降。

3. 系统级优化兼容性

TSPD 和 CE 的组合不仅减少了不必要的去噪迭代次数，还保持了输出质量。此外，这两个模块可以干净地（cleanly）与现有的系统级优化技术相结合，例如 KV Cache（键值缓存），从而在工程实现上进一步降低推理延迟。

关键要点

• 问题重构：将扩散解码从单纯的生成过程重构为动态控制问题，利用 token 级的去噪轨迹作为控制信号。
• TSPD 机制：通过轻量级控制器，结合置信度、熵、动量及位置信息，动态判断 token 收敛状态，实现并行固定。
• CE 机制：一种无需训练的模块，通过预测未来 logits 趋势和不确定性，解决轨迹震荡问题，支持安全的前瞻性决策。
• 效率提升：显著减少了冗余的去噪迭代步骤，同时未牺牲生成质量。
• 通用性：该方法对提示词和任务变化的鲁棒性优于传统的局部启发式方法，且兼容 KV Cache 等现有加速技术。
• 无额外训练成本：CE 模块是 training-free 的，意味着可以直接应用于预训练好的 dLLM，无需重新训练模型。

意义与影响

这项研究为加速基于扩散的大语言模型提供了新的技术路径。以往 dLLM 的并行生成优势常被其高昂的迭代成本所抵消，而本文提出的方法通过精细化的动态控制，有效解决了“冗余去噪”这一核心痛点。

1. 推动 dLLM 的实用化：通过降低推理延迟，使得 dLLM 在需要低延迟响应的实际应用场景中更具竞争力，缩小了其与自回归模型在速度上的差距。
2. 方法论创新：将控制理论引入文本生成过程，特别是利用“轨迹”而非单一的“步骤置信度”进行决策，为后续研究提供了新的思路。
3. 工程友好：提出的模块无需额外训练且兼容现有缓存机制，降低了落地门槛，便于集成到现有的推理引擎中。

总之，该工作通过时空并行解码与置信度外推，实现了效率与质量的平衡，是扩散语言模型推理优化领域的一项重要进展。