PSD: 通过并行推测解码推动扩散大语言模型的帕累托前沿

# PSD：通过并行推测解码推动扩散LLM的帕累托前沿 来源：https://arxiv.org/html/2605.15609 沈印欣 计算机科学系 香港城市大学 shengysun4\-c@my\.cityu\.edu\.hk &李一鸣11footnotemark:1 华为技术有限公司 li\.yiming3@huawei\.com &刘仁熙 华为技术有限公司 &李新琪 华为技术有限公司 &甄慧玲 华为技术有限公司 &林伟哲 华为技术有限公司 &陈晨 华为技术有限公司 &余贤志 华为技术有限公司 &袁明轩 华为技术有限公司 &马辰 计算机科学系 香港城市大学 chenma@cityu\.edu\.hk ###### 摘要 扩散大型语言模型通过迭代去噪掩码标记序列来生成文本。尽管扩散LLM可以在每一步内并行预测所有掩码位置，但大量的去噪迭代仍使得推理成本高昂。通过每步解掩多个标记（空间维度），或将多个去噪步骤压缩为一次验证调用（时间维度），可以降低此成本。我们提出并行推测解码——一个免训练框架，沿两个轴同时改进推理。利用单次前向传播的置信度分数，PSD通过可配置的自适应解掩策略选择要解掩的位置，并在不增加模型调用的情况下构建多深度推测草稿。最后的批量验证步骤应用层次化接受机制，保留与更新预测一致的最深草稿。在三个扩散LLM上的推理与代码生成任务实验表明，PSD在推理效率与生成质量之间实现了有利的权衡，每个前向传播的标记数最高可达5.5倍，且准确性与贪心解码相当。

## 1 引言
大型语言模型在推理和生成方面展现出卓越能力，然而自回归范式受限于逐标记顺序解码，导致高延迟和硬件利用率低下。扩散语言模型最近成为一种有前途的替代方案，将文本生成建模为对离散掩码标记序列的迭代去噪。从一个完全掩码的序列开始，扩散LLM在多个去噪步骤中逐步揭示标记，并且在每一步，模型同时为所有掩码位置生成预测。这一特性使扩散LLM天然适合并行标记生成，为更高效的推理开辟了道路。尽管有这种潜力，扩散LLM推理仍然需要多次去噪迭代，每次迭代涉及一次模型前向传播，部署效率仍然是一个实际瓶颈。

目前针对扩散LLM提出了两类加速策略。并行解码方法通过每步解掩多个标记来压缩空间维度，但随着并行度增加，质量下降，因为低置信度标记被迫过早确定，错误会传播到后续步骤。推测解码方法则通过预测未来的去噪迭代并在一次批量前向传播中验证它们来压缩时间维度；然而，现有方法每步仍只解掩一个标记，未利用空间维度，从而限制了加速效果。

然而，每种策略单独应用时都面临固有的上限。对于并行解码，增加每步解掩的标记数会减少总迭代次数，但不可避免地降低生成质量，产生任何传输策略都无法逃脱的质量-速度帕累托前沿。对于推测解码，加速受限于推测标记的接受率；由于每一步仍然只揭示一个标记，即使是完美的推测也只能带来有限的增益。因此，两个家族都在中等加速水平上饱和，留下了相当大的改进空间。

然而，我们观察到这两个加速轴在很大程度上是互补的：空间轴决定一个去噪步骤内揭示多少标记，而时间轴决定有多少个去噪步骤被压缩到一次验证调用中。我们经验性地发现，它们的加速效果结合在一起，比单独使用任何一个轴都能带来更大的收益。在这项工作中，我们提出并行推测解码，一个利用这种互补性，为扩散LLM统一空间和时间加速的框架。PSD不是规定特定的并行解码策略，而是一个通用框架，可以建立在任意并行解码算法之上。PSD基于一个关键观察：高置信度掩码位置的集合随着去噪过程的推进趋于稳定，因此在揭示更多标记后，当前被识别为高置信度的位置在很大程度上仍将是高置信度的。更具体地说，PSD分三个阶段运行：空间并行解掩应用可配置的传输策略，每步揭示多个标记；时间推测草稿通过按置信度降序填充剩余位置，在多个深度上构建候选的未来去噪结果，无需额外模型调用；以及具有层次化接受的批量验证，在单次前向传播中评估所有候选，并保留其推测标记与验证器更新预测一致的最深分支，防止激进并行解掩导致的错误传播。这些设计使得PSD能够在单次解码过程中结合两个加速家族的收益。

我们在跨越不同训练范式的三个开源扩散LLM上评估PSD，涵盖多个基准测试。经验上，PSD在所有测试配置中都实现了有利的效率-质量权衡：它通过空间并行性在仅时间方法之上提供了额外的加速，每个前向传播产生显著更多的标记，且质量损失可忽略，而同等并行度的纯并行基线则遭受明显的准确率下降。在所有模型×基准设置中，PSD始终匹配或改进了由任一家族基线建立的帕累托前沿。我们的贡献总结如下：
- •我们识别了扩散LLM中空间和时间加速之间的互补关系：这两个轴压缩了解码过程的不同维度，并且我们经验性地证明，将它们结合起来比单独使用任何一个轴都能产生更大的加速。
- •我们提出了PSD，一个免训练、策略无关的框架，在单次解码过程中将任意并行解掩策略与推测验证耦合，使用层次化接受机制来防止激进并行解掩导致的错误传播。
- •在三个扩散LLM上的数学推理和代码生成基准测试实验表明，PSD每个前向传播的标记数最高可达5.5倍，且准确率损失可忽略，始终匹配或改进了先前方法的帕累托前沿。

## 2 背景
我们的框架建立在两条工作线上：扩散LLM将文本生成建模为对掩码序列的迭代去噪，以及推测解码通过草稿-然后-验证范式加速推理。我们分别回顾如下，以激发PSD的设计。

### 2.1 扩散语言模型
作为一个新兴范式，扩散LLM将文本生成建模为对标记级掩码过程的学习逆转。在训练期间，干净的序列通过独立地用掩码标记替换标记而被破坏，模型被训练从破坏的输入中恢复原始标记。在推理时，生成始于一个完全掩码的序列并迭代进行：在每个去噪步骤，模型预测所有掩码位置的标记，然后解掩其中的一个子集，逐步细化序列直到所有位置都被填充。推理成本由k决定，即每步解掩的标记数。当k=1时，生成N个标记需要N次顺序前向传播，这使得扩散LLM处于与类似大小的自回归模型相当的延迟状态。设置k>1将总迭代次数减少到⌈N/k⌉，但迫使低置信度位置过早确定，导致明显的质量下降。这种吞吐量和准确性之间的根本张力促使我们寻找能够在不牺牲生成保真度的情况下减少前向传播次数的加速方法。

### 2.2 推测解码
推测解码通过草稿-然后-验证范式加速自回归LLM推理。不是为每个解码步骤调用大型模型，而是轻量级的草稿模型首先提出几个未来的标记，然后大型目标模型在单次前向传播中评估这些提议。其草稿预测与目标模型一致的标记被接受，因此一次昂贵的目标模型调用可以推进解码过程多个步骤，同时仍然保留目标分布。然而，将这种范式适应扩散LLM是非平凡的。在自回归生成中，标记以固定的从左到右顺序产生，因此下一个预测目标总是明确的。相比之下，在扩散LLM中，解掩顺序由模型置信度动态确定，且事先未知；因此，推测的单位从单个下一个标记转变为要在每一步揭示的整个位置集。

参见图注
图1：提出的PSD概述。作为一个与任何并行解码策略兼容的通用框架，PSD包含三个阶段：(1) 空间并行解掩在每个去噪步骤揭示多个标记；(2) 时间推测草稿为未来步骤构建置信度排序的草稿有向无环图；(3) 验证与接受在一次批量前向传播中评估所有草稿，并接受最深的一致分支。通过结合空间和时间加速，PSD实现了超越单独任何一个轴的复合速度提升。

## 3 方法
在本节中，我们首先形式化扩散LLM解码过程，然后介绍PSD的三个阶段：空间并行解掩通过任意传输策略每步揭示多个标记，时间推测草稿构建候选的未来结果，以及具有层次化接受的批量验证在单次前向传播中保留最深的一致分支。

### 3.1 预备知识
将去噪步骤t的序列状态记为x\(t\)∈\(V∪\{\[M\]\}\)N，其中V是词汇表，\[M\]是特殊的掩码标记。在每个步骤t和每个掩码位置i，模型为每个标记v∈V分配概率pθ\(xi=v∣x\(t\)\)。我们将最可能的标记记为x^i\(t\)=argmaxv∈V pθ\(xi=v∣x\(t\)\)，其置信度为ci\(t\)=maxv∈V pθ\(xi=v∣x\(t\)\)。令M\(x\(t\)\)={i:xi\(t\)=\[M\]\}为步骤t时掩码位置的集合。传输策略π然后选择在该步骤解掩的位置子集T\(t\)⊆M\(x\(t\)\)，并且r̄π=E[|T\(t\)|]表示每步揭示的平均位置数。

### 3.2 并行推测解码
如图1所示，PSD将空间和时间加速整合到一个统一的三阶段流水线中。
阶段 (1) 空间并行解掩通过*任意*传输策略在每个去噪步骤揭示多个标记，PSD对此选择不可知。阶段 (2) 时间推测草稿重用当前步骤的置信度分数，在不增加模型调用的情况下构建候选的未来去噪结果的有向无环图。阶段 (3) 具有层次化接受的批量验证在单次前向传播中评估所有草稿候选，并接受与验证器更新预测保持一致的最深分支。下面我们详细介绍每个阶段。

#### 空间并行解掩。
尽管诸如置信度阈值之类的传输策略原则上可以在整个序列上操作，但最先进的扩散LLM将生成组织成连续块，长度为L，这是包括LocalLeap和LoPA在内的近期方法共享的设计。在块边界处应用PSD是直接的：将块视为基本单元，并行策略决定每个块内解掩的标记数。设r表示目标每步解掩率。然后，在步骤t，并行策略选择具有最高置信度ci\(t\)的至多r个掩码位置。然而，由于扩散LLM块架构需要完整的块上下文才能进行有效预测，直接对整个序列进行操作可能会破坏块级一致性。因此，我们在每个块内独立地应用此选择：对于每个长度为L的块，我们选择该块内置信度最高的min(r, |M∩块|)个位置。这种块内并行解码确保了每个块在解掩步骤中接收一致的上下文，同时保持所需的解掩率。

#### 时间推测草稿。
一个关键的观察是，置信度分数ci\(t\)在相邻的去噪步骤之间保持稳定：在某一时刻被模型高度确定的位置，在揭示几个附近的标记后，通常仍然保持高度确定。基于此，我们通过从当前置信度分数直接推断构建候选的未来序列。具体来说，对于每个掩码位置i，我们假设其预测标记x̂i\(t\)在未来步骤中保持不变，其置信度ci\(t\)作为在未来步骤中它是否将被解掩的代理指标。这允许我们生成多个依次减少掩码的草稿。对于深度d，我们根据当前置信度对剩余掩码位置进行排序，并解掩前d·r个位置（其中r是每步解掩率）。将此应用于从步骤t开始的每个未来深度，产生一个深度树，其中每个节点对应一个部分解掩的序列。至关重要的是，所有这些草稿都是仅基于步骤t的置信度分数在计算上构建的，无需任何额外的模型前向传播。

#### 批量验证与层次化接受。
最后，在验证阶段，我们使用目标扩散LLM评估所有构建的草稿。考虑一个深度为D的草稿树。对于每个深度d∈{1,…,D}，我们通过将相应的草稿序列与当前步骤t的序列进行比较来定义一致性度量。如果对于所有标记j，草稿的解掩标记与模型在步骤t+δ的预测（即目标模型在进行δ步去噪后的输出）相匹配，则草稿d被认为是“一致的”，其中δ由树的深度和步长决定。由于目标模型通常不会提前输出，我们必须考虑草稿与真实模型更新之间的分歧。我们的层次化接受机制操作如下：(1) 在单次前向传播中，我们计算目标模型在所有草稿深度上的概率。(2) 我们从最深的草稿（最高D）开始，如果草稿d在整体迭代的一致性标准下与目标预测一致，我们接受它，并相应地更新序列。(3) 如果不一致，我们移至下一最浅的草稿，并重复此过程，直到找到一致的草稿或达到基本步骤（仅接受空间并行解掩的输出）。此过程通过要求所有更深的草稿完全匹配目标输出来防止错误积累——如果草稿d因任何差异而被拒绝，这确保我们不会因接受冲突的草稿而向后传递错误。通过一次验证调用，我们有效地确认了多个去噪步骤的准确性，同时通过层次化门控机制防止低置信度决定的传播。