LiLiCorr：用于推测解码的并行草稿轻量似然相关性

# LiLiCorr：用于推测解码的轻量级并行草稿似然关联  
来源：https://arxiv.org/html/2608.20530  
Matan Rusanovsky、Yoav Miron、Roy Uziel、Omer Belhasin、Ran Zilberstein、Maor Ashkenazi、Michael Elad  
\[2\pt\] NVIDIA  

###### 摘要  
推测解码通过起草未来令牌并由目标模型并行验证的方式加速语言模型推理。类似 DFlash 的扩散式块头是一个有吸引力的起草器，它在单次前向传播中预测整个未来令牌块。然而，它是基于逐位置边缘分布而非联合块分布训练的，因此其生成的令牌各自看似合理但整体不连贯。我们提出 LiLiCorr，一个轻量级基于似然的模型，用于关联起草器已生成的逐位置边缘分布。该模型在每个位置保留前-K 个令牌作为*候选*并联合处理它们，为每个候选生成一个*入向量*和一个*出向量*。这些向量链接连续位置的候选：当前驱候选的*出向量*与后续候选的*入向量*具有高余弦相似度时，两者即匹配。这些匹配捕捉了块的联合结构，而无需显式表示随块长度呈指数增长的完整联合分布。训练使正确配对得分最高，同时压低竞争配对，从而使连贯块得分高于不连贯块。一次轻量级网络前向传播生成所有向量，随后成对得分通过批处理矩阵运算并行计算，仅剩一个廉价的贪心遍历步骤是顺序执行的。我们进一步将 DFlash 起草器与 LiLiCorr 联合训练，使其学会提出能关联成长接受序列的候选。与原始 DFlash 起草器相比，LiLiCorr 在所有基准测试上将接受长度提升了 9% 至 19%，而其单次前向评分头仅占每块延迟的约 2.8%。与 DFlash 及另外两种在起草时恢复连贯性的并发方法相比，LiLiCorr 在 72 种设置中的 70 种中提供最高吞吐量：在贪心和温度为一的解码下，针对两个目标规模的九个基准测试；以及在六种并发度、两种输入长度和三种熵层级上的吞吐量扫描，所有系统在通用服务栈上均进行了同等优化。将 LiLiCorr 扩展至比其训练长度长一个数量级的输入时，该优势仍然保持。  
††\*通讯作者：Matan Rusanovsky，mrusanovsky@nvidia\.com。  

## 1 引言  
推测解码通过让一个小型草稿模型提议多个未来令牌并由目标模型并行验证的方式加速自回归语言模型推理，同时保留最长的正确前缀（8 (https://arxiv.org/html/2608.20530#bib.bib1)，3 (https://arxiv.org/html/2608.20530#bib.bib2)）。由于解码受内存带宽限制，一次性验证大量令牌的代价与仅生成一个令牌相差无几。由此产生的加速取决于两个量：*接受长度*（每步接受的草稿令牌期望数量）和草稿模型本身的成本。一个好的草稿模型必须提高前者而不增加后者。因此，大量工作致力于提高接受长度，设计范围从 EAGLE 等特征自回归头（10 (https://arxiv.org/html/2608.20530#bib.bib4)，9 (https://arxiv.org/html/2608.20530#bib.bib5)，11 (https://arxiv.org/html/2608.20530#bib.bib6)）到 Medusa 等多头预测器（2 (https://arxiv.org/html/2608.20530#bib.bib3)）。一种特别有前景的方法通过类似 DFlash 的并行扩散式草稿器在单次前向传播中预测整个未来令牌块（4 (https://arxiv.org/html/2608.20530#bib.bib7)），以极低成本提供长推测窗口。  
在本工作中，我们基于 DFlash 构建，它在一次前向传播中预测整个未来令牌块。对于块中的每个位置（后续称为*槽位*），它生成完整词汇表上的分布。尽管 DFlash 在槽位间进行双向注意力，使其表示原则上可以关联不同位置，但它通过仅监督每个槽位令牌*边缘*的逐位置交叉熵进行优化。这使得块上的*联合*分布基本不受约束，因此 DFlash 能很好地预测每个槽位，但可能将令牌组合成不连贯的块。  
这一边缘问题在离散扩散文献中众所周知，其根源是组合性的：块中 B 个位置（每个位置从大小为 V 的词汇表中抽取）上的联合分布有 V^B 种可能实现，因此即使对于非常短的块，显式表示它也是难以处理的。扩散语言模型因此从不显式建模联合分布，而扩散和推测解码文献都提供了在没有它的情况下恢复连贯性的方法。  
掩码扩散模型（如 LLaDA、MDLM 和 Dream（13 (https://arxiv.org/html/2608.20530#bib.bib14)，17 (https://arxiv.org/html/2608.20530#bib.bib15)，21 (https://arxiv.org/html/2608.20530#bib.bib16)）通过迭代去噪关联序列，在若干轮中重新掩码部分预测令牌并重新预测它们。然而，在推测解码设置中，每一轮都是一次完整的草稿模型前向传播，代价过高。  
树状重排序（如 DDTree（16 (https://arxiv.org/html/2608.20530#bib.bib10)））从逐槽位边缘扩展多条候选路径，并让目标模型从中选择。这将关联工作转移到目标模型上，而推测解码的设计正是为了节省目标模型的传递次数。这种交换在计算受限的高并发现实服务场景中尤为不利。  
相比之下，Domino（7 (https://arxiv.org/html/2608.20530#bib.bib8)）和 DSpark（5 (https://arxiv.org/html/2608.20530#bib.bib9)）这两个与我们同步开发的工作在起草时采取行动，避免额外的目标模型传递，但它们是通过一系列自回归神经网络传递实现的，在每个草稿槽位上运行一次。  

DFlash：一次前向传播  
词汇表上每个槽位的边缘分布  
槽位1 槽位2 槽位3  
每个槽位的 argmax  
The dog meows sweetly ×\times “狗温柔地喵喵叫”  
每个槽位的 top-K 候选  
dog cat bird meows barks sings sweetly loudly softly  

LiLiCorr：一次双向传递  
in out  
0.34 0.81 0.22  
The dog cat bird meows barks sings sweetly loudly softly  
所有链接都并行计算 ✓\checkmark “猫温柔地喵喵叫”  
一次网络传递，所有槽位同时处理  
argmax 遍历，无网络权重  

(a) 边缘起草器联合错误  
(b) 一次传递遍历所有候选  
(c) 廉价贪心解码  

图 1：LiLiCorr 概览。(a) DFlash 在一次前向传播中为每个槽位生成词汇表上的分布 logit，并独立地取每个槽位最可能的令牌，会产生局部合理但整体不连贯的序列。(b) LiLiCorr 保留每个槽位的 top-K 个候选令牌，并在一次传递中联合评分，为每个候选生成一个*入向量*和一个*出向量*。(c) 相邻候选通过前驱*出向量*与后续*入向量*的对齐进行耦合，因此每个槽位对简化为一次矩阵乘法，整个块并行评分。然后块从左到右通过贪心 argmax 操作提交，该操作不涉及任何网络权重。  

我们提出 LiLiCorr，一个轻量级基于似然的模型，能在*单次*网络传递中从这些逐槽位边缘恢复连贯序列。图 1（https://arxiv.org/html/2608.20530#S1.F1）展示了端到端的方法。LiLiCorr 接收每个槽位的 top-K 个候选令牌，并为每个候选生成两个向量：*入向量*和*出向量*。这些向量捕获候选如何链接在一起：当两个相邻候选兼容时，即前驱候选的*出向量*与后续候选的*入向量*在余弦相似度上对齐。训练使连续真实令牌的*出向量*和*入向量*对齐，并分离竞争配对。在推理时，廉价的贪心遍历从左到右提交块，在每个位置选择与其已提交前驱最匹配的候选。  
核心新颖之处在于，这种关联逻辑简化为*易于优化的矩阵操作*。所有候选交互都源于一次前向传播并并行计算。只有贪心路径选择是顺序的，这与之前每个槽位运行一次关联网络的起草时方法不同。LiLiCorr 与 DFlash 联合训练，而非在冻结的起草器之上训练，因此起草器学会提议 LiLiCorr 能转化为更长接受序列的候选集。  
与在同一数据上训练的原始 DFlash 起草器相比，这在所有基准测试上将接受长度提高了 9% 至 19%，适用于 8B 和 4B 目标规模以及贪心和温度为一的解码。我们进一步与上述两个并发工作 Domino 和 DSpark 进行比较，它们同样与其头模块一起调整起草器，让所有系统在通用服务栈上获得相同的优化努力。LiLiCorr 在 83% 的基准测试上比 Domino 实现更长的接受长度，并在所有设置的基准平均值上领先；它比两个基线都更高效地将草稿转化为令牌：在所有三十六种设置中，它比两者服务都快，在基准平均值上领先较强基线多达 5.3%。原因是结构性的：两个基线都在每个草稿槽位上花费一次网络传递，而 LiLiCorr 一次性对整个候选网格评分。这导致其头模块比 Domino 便宜约 2.4 倍（两者完全优化后）（表 3（https://arxiv.org/html/2608.20530#A1.T3)）。  
随着服务并发度上升，该优势持续存在，并且在推理时扩展位置编码后，即使输入比草稿器训练长度长一个数量级，该优势依然保持。  

##### 贡献。  
1. **(1) 一种新颖的基于似然的边缘问题公式化**。我们将从逐槽位边缘恢复连贯序列的过程构建为候选路径上的学习评分模型。模型不是在整个词汇表上推理，而是通过为每个槽位的所有候选分配两个低维向量（编码邻居如何匹配）来关注所有候选。  
2. **(2) 关联作为高效的并行矩阵操作**。在解码时，一次网络前向传播预计算所有关联，因为成对候选交互简化为易于优化的矩阵操作。随后是一个廉价的贪心遍历，该步骤保持顺序执行。而先前的起草时方法则在每个槽位上运行一次关联网络。  
3. **(3) 最先进的结果**。LiLiCorr 在推测解码的起草时关联方面设定了新的最先进水平，它与起草器联合训练，并获得了与两个并发工作（7 (https://arxiv.org/html/2608.20530#bib.bib8)，5 (https://arxiv.org/html/2608.20530#bib.bib9)）相同的数据、训练预算和优化努力。我们在九个基准测试、两个目标规模、贪心和温度为一的解码下进行了评估，并在六种服务并发度、两种输入长度和三种输出熵层级上进行了吞吐量扫描。在 72 种设置中的 70 种中，LiLiCorr 保持任何被比较系统中的最高吞吐量，并且该优势在输入长度超过草稿器训练长度时仍然存在。  

## 2 相关工作  
##### 推测解码与并行草稿器。推测解码通过提议草稿令牌并在保持其分布规则下使用目标模型验证它们来加速推理（8 (https://arxiv.org/html/2608.20530#bib.bib1)，3 (https://arxiv.org/html/2608.20530#bib.bib2)）。大多数此类工作改进起草器，从 Medusa 等多头预测器（2 (https://arxiv.org/html/2608.20530#bib.bib3)）到 EAGLE 家族等特征自回归头（10 (https://arxiv.org/html/2608.20530#bib.bib4)，9 (https://arxiv.org/html/2608.20530#bib.bib5)，11 (https://arxiv.org/html/2608.20530#bib.bib6)）。另一条工作线在一次前向传播中起草整个块：并行和扩散式块头（如 DFlash（4 (https://arxiv.org/html/2608.20530#bib.bib7)）和基于扩散的自我推测（6 (https://arxiv.org/html/2608.20530#bib.bib13)））能廉价处理长推测窗口。D-PACE（19 (https://arxiv.org/html/2608.20530#bib.bib11)）通过接受感知训练目标改进了 DFlash。它们的并行预测类似于扩散语言模型，如 LLaDA（13 (https://arxiv.org/html/2608.20530#bib.bib14)）、MDLM（17 (https://arxiv.org/html/2608.20530#bib.bib15)）、Dream（21 (https://arxiv.org/html/2608.20530#bib.bib16)）和 Nemotron-Labs-Diffusion（6 (https://arxiv.org/html/2608.20530#bib.bib13)）。这些模型表现出众所周知的边缘限制：尽管注意力是双向的，但块被分解为独立的逐位置预测，因此采样令牌不能保证联合连贯。这通过迭代细化来解决，重新掩码并重新预测令牌子集，以便后续轮次以早期轮次为条件。并行草稿器继承了相同的限制，但在推测解码中，每轮细化都增加一次完整的草稿模型传递，这抵消了并行起草所激发的延迟收益。LiLiCorr 解决的问题是：如何在不支付重复草稿模型传递代价的情况下，从这些边缘恢复连贯块。  
##### 恢复连贯草稿。几种最近的方法从并行草稿器（如 DFlash）产生的逐槽位边缘恢复连贯序列。一个方向将这些边缘扩展为草稿树，包含多条目标在单次传递中验证的延续路径。DDTree（16 (https://arxiv.org/html/2608.20530#bib.bib10)）无需训练，通过边缘乘积对候选前缀进行排名，而 14（https://arxiv.org/html/2608.20530#bib.bib12) 先学习一个小的自回归适配器再构建树。两者都将关联工作转移到目标模型上，花费了推测解码旨在节省的计算资源，并且都依赖树状注意力，这通常比标准因果注意力效率低。  
另一个方向在起草时重新引入依赖关系：Domino（7 (https://arxiv.org/html/2608.20530#bib.bib8)）和 DSpark（5 (https://arxiv.org/html/2608.20530#bib.bib9)），这两个与我们同步的工作，使每个槽位以之前采样的令牌为条件，这避免了额外的目标模型传递，但在每个槽位上运行一次校正网络。两者也像我们一样联合调整起草器及其校正头。LiLiCorr 也在起草时采取行动，但仅需一次网络传递：它一次生成所有候选表示，并行计算其成对关联，并仅剩一个廉价的贪心 argmax 操作在生成的分数上顺序运行。  

## 3 方法  
### 3\.1 设置：前缀接受与边缘起草器  
设 M 为冻结的目标语言模型。一个并行块起草器在一次前向传播中提议一个 B 个位置的块，该块基于已验证的上下文，其中 B−1 个是未来令牌。块的第一个位置是目标已提交的最后一个令牌，作为*锚点*，槽位 0。剩余的 B−1 个位置是推测*槽位* 1,...,B−1。在每个槽位，起草器生成词汇表上的分布（logit），即逐槽位*边缘*。标准解码独立解决每个槽位，取其最可能的令牌。草稿完成后，目标在一次传递中验证整个块，并接受其一个*前缀*，丢弃从第一个被拒绝位置开始的所有内容。对于贪心目标，草稿令牌在与其自身选择一致时被接受。在采样目标下，接受则遵循标准规则。