基于时空并行解码与置信度外推的高效扩散LLMs

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

来源：https://arxiv.org/html/2605.30753

###### 摘要

基于扩散的大语言模型（dLLMs）通过迭代去噪支持并行文本生成，然而其推理过程仍存在高延迟问题，因为许多步骤被用于对已经确定最终值的 tokens 进行冗余优化和重复掩码。先前的加速方法主要依赖于基于单步置信度的启发式规则或固定调度，这些方法对提示词和任务的变化较为敏感，且忽略了序列内部显著的位置效应。本文将扩散解码视为一个动态控制问题，并揭示出 token 级的去噪轨迹为可靠的决策控制提供了关键信号。我们提出了一种轨迹感知的解码框架，包含两个组成部分。首先，**时空并行解码（Temporal-Spatial Parallel Decoding, TSPD）** 使用一个轻量级的时空控制器，它利用每个 token 的轨迹特征（包括置信度、熵和动量）以及 token 位置，来判断一个 token 何时已收敛并可以安全固定。其次，我们引入了**置信度外推（Confidence Extrapolation, CE）**，这是一个无需训练的状态空间模块，用于预测未来 logits 的趋势及其不确定性，以支持主动决策，包括在轨迹振荡或置信度不足时进行安全的提前预测和目标性稳定化。TSPD 和 CE 共同减少了不必要的去噪迭代，同时保持了输出质量，并且能够与 KV 缓存等系统优化手段无缝结合。

机器学习，ICML

## 1 引言

基于扩散的大语言模型（dLLMs）[^1] 提供了一种自回归解码的替代方案，通过迭代去噪并行更新多个 token 位置。这种框架支持灵活的生成顺序、文本填充以及通过去噪步数直接调节的质量-速度权衡[^2]。近期出现的掩码扩散语言模型[^3] 和 LLaDA 风格模型[^1] 表明，扩散模型在规模上可以具备竞争力。尽管具有并行性，但在实践中推理延迟仍然很高。每个去噪步骤通常需要在双向注意力机制下对整个序列进行一次前向传播[^4]，并且很多步骤被用于优化最终值已经确定的 token。当不确定性集中在少量位置上时，重复掩码会进一步放大这种成本。

先前的研究通过复用跨步骤的计算、调整解掩码调度或在每步承诺更多 token 来减少这种开销[^5]。这些方法提高了吞吐量，但大多数提交和停止决策仍然依赖于单步启发式规则、固定调度或全局阈值[^6]。这类信号较为脆弱，因为扩散模型的置信度在不同提示词、领域和 token 位置上并不一致。实验表明，扩散解码表现出结构化的时间和空间行为，包括延迟稳定、置信度不足但却正确的 token，以及振荡的置信度[^7]。单步快照或启发式规则无法区分这些状态。因此，局部规则常常在收敛后仍进行冗余去噪，或者过早提交而损失质量。

在本工作中，我们将扩散解码视为一个**序列控制问题**，而非一系列独立的阈值测试。一个关键的观察是，去噪过程在时间步上产生了 token 级的*轨迹*，例如置信度和熵的曲线，其时间模式同时揭示了正确性和稳定性。这使我们从关注“token 在第 t 步是否自信”转向关注“轨迹是否指示收敛，以及如果继续去噪，它将如何变化”。回答这个问题需要一个时空感知的控制器，它能够解释诸如稳定增长、振荡和延迟稳定等模式，以及一种能够利用可预测趋势来避免不必要步骤的机制。

我们提出了**时空并行解码（Temporal-Spatial Parallel Decoding, TSPD）**，这是一个时空感知框架，旨在提升效率和准确性。TSPD 引入了一个轻量级的**时空控制器**，利用序列化的轨迹特征（最近几步的置信度、熵和动量）以及块内的相对空间位置，来预测某个位置是否已准备好被固定。一个共享的序列感知模型捕捉了单步分类器和手工调优阈值经常遗漏的动态特性，从而提高了跨提示词和位置效应的鲁棒性。为了减少当置信度上升但尚未决定性时的等待时间，我们进一步引入了一个无需训练、即插即用的**置信度外推机制（Confidence Extrapolation, CE）**。CE 将置信度演化建模为一个简单的状态空间过程，并预测未来的置信度及其不确定性。一个风险感知的视界规则根据左侧上下文进展和预测不确定性来限制提前预测范围，确保外推仅在可靠历史支持的情况下应用。

我们在 LLaDA-8B-Instruct[^1] 上评估了我们的框架（TSPD 和 CE），涵盖了算术推理和代码生成四个基准：GSM8K[^8]、MATH[^9]、HumanEval[^10] 和 MBPP[^11]。在多个 dLLM 骨干网络和生成设置下，我们的方法一致地提供了比现有最先进基线更强的速度-质量权衡。在 GSM8K 上生成 256 个 token 时，它在无缓存情况下实现了 **5.0×** 加速，与 KV 缓存结合后加速比达到 **11.2×**，表明我们控制器层的增益与系统优化是互补的。优势在生成长度增加时更加明显：生成 1024 个 token 时，我们的方法达到了 64.1 Tokens Per Second (TPS)，加速比高达 **58.3×**，同时准确率基本保持不变，在某些情况下甚至略有提升。总结而言，我们的贡献有三点：

- 我们提出了 **TSPD**，一个时空感知的解码控制器，利用时间置信度动态和位置上下文进行稳健的 token 固定决策，缓解了单步启发式规则的脆弱性。
- 我们引入了一个**无需训练、即插即用的置信度外推模块**，该模块在状态空间模型下预测未来置信度，并量化不确定性，同时应用风险感知的提前预测来补充多样化的解码控制器。
- 我们在广泛的 dLLMs、任务领域和生成设置上进行了大量评估和消融实验，证明了显著的加速效果，同时准确率几乎无损失，并且与 KV 缓存完全兼容。

## 2 相关工作

### 2.1 基于扩散的大语言模型

扩散模型[^12] 已从连续数据扩展到离散序列，通过迭代去噪实现了非自回归生成。早期的离散扩散工作建立了类别空间上的马尔可夫公式，包括多项扩散[^13] 和 D3PM[^14]。连续时间视角[^15] 进一步统一了离散扩散，并支持灵活的采样调度。SEDD[^16] 通过使用分数-熵目标对边际概率的比率进行建模，改进了学习过程，缩小了与自回归语言模型（LMs）的差距。

基于这些基础，掩码扩散大语言模型已成为文本生成的一种实用范式。MDLM[^17] 表明，简化的掩码扩散方案可以显著提高困惑度和采样质量。在大规模上，扩散大语言模型要么通过改编强大的自回归骨干网络（如 DiffuLLaMA）[^18]，要么从头开始训练（如 LLaDA）[^1] 来开发，在指令遵循和推理方面取得了有竞争力的表现。Dream 7B[^19] 进一步推进了开放式扩散大语言模型，并突出了扩散模型特有的能力，如任意顺序生成和文本填充，同时具备可调节的质量-速度权衡。

参见图注图 1：相同的局部置信度，不同的轨迹可能导致不同的解码结果。
参见图注图 2：空间位置越靠右的 token，稳定时间越晚。

### 2.2 扩散大语言模型的解码加速

尽管具有并行更新，但扩散解码的延迟仍然可能很高，因为每个去噪步骤通常需要在双向注意力机制下对整个序列进行一次前向传播。先前的工作主要通过三个方向来提高效率：(i) 计算复用和缓存，(ii) 更好的并行提交或采样调度，以及 (iii) 即插即用的加速器。

在计算复用方面，dLLM-Cache[^20] 缓存静态提示词并选择性更新稳定的 token；Fast-dLLM[^21] 引入了块级近似 KV 缓存，结合置信度感知的并行解码；DInfer[^22] 将推理流水线模块化，并报告相对于 Fast-dLLM 实现了约 10 倍的加速且无质量损失；Fast-dLLM v2[^23] 进一步支持将预训练的自回归骨干网络进行块扩散适配，以及分层 KV 缓存，实现了相对于标准自回归解码高达 2.5 倍的加速。在算法解码方面，SlowFast Sampling[^24] 交替进行保守探索和加速解码；Local Leap[^25] 利用高置信度锚点周围的局部确定性；Learn2PD[^26] 学习一个并行解码策略，在 LLaDA 上实现了显著加速；自适应并行解码[^27] 动态控制并行度，并使用一个小型自回归验证器来实现可控的速度-质量权衡。无需训练的即插即用插件包括 FlashDLM[^28]（FreeCache + 引导扩散）、ES-dLLM[^29]（通过收敛信号早期跳过）和预见性移动解码[^30]。

现有的大多数方法仍然基于单步启发式规则或固定调度做出决策，并通过被动地运行更多去噪步骤来应对不确定性。相比之下，我们强调扩散解码揭示了信息丰富的 **token 级轨迹**（置信度/概率曲线），从而能够实现更具输入适应性的控制。因此，我们开发了 (i) 一个**时空正确性传感器**，融合历史特征与位置线索以推断稳定状态，以及 (ii) 一个**主动预测模块**，该模块使用具有校准不确定性的方法（例如，通过状态空间/卡尔曼风格模型[^31]）外推近期的置信度，以便进行风险受限的提前决策。这种轨迹感知的视角补充了缓存和采样器的优化，并针对扩散大语言模型推理中更细粒度的速度-质量控制。

参见图注图 3：被动等待与提前预测相比错过的加速机会。
参见图注图 4：可被外推的一致步骤比例分布。

参见图注图 5：标准 dLLM 并行解码（左）与我们的框架（右）的比较。标准解码在每次全序列去噪传播后应用单步启发式规则，经常重新访问已经正确的 token。我们的框架插入了一个风险感知的置信度外推器（CE）和一个时空解码控制器（TSPD）：CE 提供带有不确定性感知的提前置信度预测，TSPD 使用 token 级别的置信度轨迹和位置上下文来固定稳定的 token。这减少了冗余步骤，同时保持了生成质量。

## 3 方法

### 3.1 预备知识

#### 3.1.1 扩散大语言模型的并行解码

扩散大语言模型（dLLMs）通过迭代去噪一个序列的掩码版本，生成一个长度为 \(L\) 的序列 \(x_0 \in \{0, 1, \ldots, V-1\}^L\)。设 \(\alpha \in [0,1]\) 表示掩码比例（噪声级别），以避免将 t 重载为“时间”。前向过程独立地将每个 token 以概率 \(\alpha\) 替换为掩码 token \(m\)：

\[
q_\alpha(x^\alpha \mid x) = \prod_{i=1}^{L} \left[ (1-\alpha) \mathbf{1}[x_i^\alpha = x_i] + \alpha \, \mathbf{1}[x_i^\alpha = m] \right].
\tag{1}
\]

其中 \(x_\alpha\) 是在掩码比例 \(\alpha\) 下被破坏的序列。逆向过程在 \(K\) 个离散步骤中去噪，通过反复预测掩码位置的 token 分布。给定一个提示词 \(c = (c_1, \ldots, c_M)\) 和当前掩码响应 \(y^{(k)}\)，一个掩码预测器 \(p_\theta\) 输出每个位置的分布，我们通过贪婪解码形成一个临时序列：

\[
\hat{y}^{(k)}_i = \arg\max_{v \in [V]} p_\theta\left(v \,\middle|\, c, y^{(k)}\right), \quad i \in [L].
\tag{2}
\]

然后对 \(n_k\) 个具有最小 \(c_i\) 的位置进行重掩码，其中 \(n_k\)（相当于掩码比例）遵循一个预定义的时间表，随着步骤递减。这

(Note: The original text cut off. The translation continues from where the English text ended. I will complete the translation of the rest of the paper section by section.)

...