复制还是不复制：通过内在模型信号控制推测解码

# 要复制还是不复制：通过内在模型信号控制推测解码
来源：https://arxiv.org/html/2609.20186

###### 摘要

推测解码 \(SD\) 已显著加速了大语言模型 \(LLM\) 的推理，然而现有方法在两种草稿策略之间面临根本性的权衡：神经草稿和基于上下文的复制。神经草稿（例如 EAGLE3）在各种文本设置下提供稳健的性能，而基于复制的方法在复制密集型场景中通过更快生成候选方案并利用长重复跨度实现近乎完美的推测，从而获得更高的加速比。我们分析了现有的基于复制的方法，发现它们容易产生**偶然重复**，即表面的n-gram重叠并未反映出结构化的复制意图，导致误报触发，最终降低吞吐量。我们引入了 SwitchSD，这是一个自适应框架，将复制视为 LLM 的一种**潜在控制信号**。通过对目标模型的内部表示训练轻量级探测器，SwitchSD 以高精度（AUC > 0.99）识别**真正的复制意图**。这允许系统在神经草稿（例如 EAGLE）和基于上下文的复制之间动态切换。我们在 Llama 和 Qwen 系列模型上的结果表明，与 EAGLE3 等最先进基线相比，吞吐量提升了高达 15%，有效地将“复制”从一种嘈杂的启发式方法转变为一种有原则的、模型感知的解码模式。

††脚注：1特拉维夫大学 2特拉维夫隐形初创公司

## 1 引言

大语言模型 \(LLMs\) 重新定义了几乎所有自然语言处理领域的尖端性能，常常表现出涌现的推理能力。然而，这些增益越来越受到自回归解码计算成本的限制。这一瓶颈在**测试时扩展**场景中最为突出，因为生成详尽的思维链 \(CoT\) 序列会显著增加延迟。随着 LLM 向更深层次的推理发展，解码的顺序特性仍然是实际部署的主要障碍。

推测解码 \(SD\) 已成为解决这一瓶颈的标准方法，它利用一个轻量级的草稿模型来提议令牌序列，然后由目标模型并行验证。除了标准的神经草稿，非神经方法如 **Prompt Lookup Decoding \(PLD\)** 和基于检索的推测器试图通过 n-gram 搜索直接从上下文中复制来绕过神经开销。基于这些思想，最近的创新如 **CopySpec** 引入了一种混合启发式方法：如果前缀先前出现在上下文中，系统就推测后续的延续；如果未找到先前的出现，则回退到神经草稿策略。虽然这种方法很直观，但它依赖于一个脆弱的假设，即表面的 n-gram 重叠总是意味着存在复制上下文的结构意图。

如图1所示（推测解码过程示意）：在目标模型的每次验证步骤后，给定最终接受的令牌（以绿色表示），存在三种可能的情况：(i) 没有包含当前生成令牌的重复。在这种情况下，使用标准的神经草稿模型。否则，存在包含当前令牌的重复。在这种情况下：(ii) 如果探测器（以红色表示）识别出复制意图，则通过从上下文复制来起草推测令牌；(iii) 如果未识别出复制意图，则使用神经草稿模型。神经草稿模型以灰色显示，复制草稿以橙色显示，生成的令牌以蓝色显示，重复的输入令牌带下划线。图中的示例按其编号顺序包含了这三种情况。

我们的分析揭示了这些表面启发式方法的一个关键缺陷：它们无法区分**真正的复制意图**和**偶然重复**。如图2所示，重复的前缀往往源于巧合的联想模式（例如，在数学推理中），而不是对结构化文本的刻意复制。在这些偶然情况下触发基于复制的草稿会导致立即拒绝，产生“推测税”，从而降低吞吐量，不如标准的神经草稿。我们认为，复制不是一种表面现象，而是由模型内部表示中编码的**潜在控制信号**驱动的。

受此启发，我们提出了 **SwitchSD**，这是第一个模型感知的 SD 框架，它用一个有原则的、基于表示的决策机制取代了嘈杂的启发式方法。通过直接从目标模型的**内在表示**中提取潜在的复制信号，SwitchSD 能够高精度地识别复制意图，从而实现动态编排：仅当模型处于“复制模式”时通过上下文复制进行推测，否则回退到最先进的神经推测器（例如 EAGLE3）。这种选择性机制消除了误报的复制尝试，将上下文复制从一种脆弱的启发式方法转变为一种稳健、高收益的加速策略。

**情况 A：潜在复制意图（高收益）**
上下文（来源）：
def truncate(n):
    """给定一个正数，它可以被分解为整数部分和小数部分..."""
当前前缀：
Assistant: def truncate(number):
目标（有意复制）：
"""给定一个正数，它可以被分解为..."""
复制草稿："""给定一个正数，它可以..."""
结果：20个令牌被接受
神经草稿："""\n"""
结果：2个令牌被接受

**情况 B：偶然重复（零收益）**
上下文（问题）：
三个连续整数的最小完美立方和是...
...接下来，让我们尝试 n=2。如果 n=2，和是 9（不是立方数）。
当前前缀：
Next, let's try n =
目标（新颖推理）：
3。如果 n=3，那么和是12...
复制草稿（天真）：2。如果 n=2，那么...
结果：0个令牌被接受
神经草稿：3。如果 n=3，那么...
结果：8个令牌被接受

图2：潜在复制意图与偶然重复。(左) SwitchSD 在文档字符串中识别出复制意图。(右) 在推理任务中，n-gram 重叠是偶然的；SwitchSD 正确地回退到神经草稿。

**我们的主要贡献**如下：(i) 我们引入了 SwitchSD，这是一个自适应 SD 框架，通过对目标模型隐藏表示进行轻量级探测，在神经草稿和上下文复制之间动态切换。我们的方法与现有 SD 方法正交，并在多个基准测试和模型系列中一致地改进了包括 EAGLE3 在内的强大基线，吞吐量提升高达 15%。(ii) 我们表明 SD 包含异构的解码模式，其接受长度特征存在显著差异。特别是，复制有利的区域产生的接受续写比标准生成区域显著更长。利用这种结构可以实现更有效的草稿策略选择和改进的推测前瞻调整。(iii) 通过广泛的消融实验和分析，我们证明了内部模型表示为自适应 SD 提供了可靠的控制信号，并系统地分析了探测器设计、训练数据、阈值选择和草稿策略对整体解码效率的影响。

## 2 背景与相关工作

本节为 SD 和上下文复制机制建立正式框架。

##### 自回归 Transformer。
现代 LLMs 由堆叠的 Transformer 块组成，其中隐藏表示 \(H^{\ell}\) 通过自注意力和 MLP 子层进行转换。在预归一化公式下，一个块定义为：
\[
\tilde{H}^{\ell} = H^{\ell} + \mathrm{Attn}\!\left(\mathrm{LN}(H^{\ell})\right), \qquad H^{\ell+1} = \tilde{H}^{\ell} + \mathrm{FFN}\!\left(\mathrm{LN}(\tilde{H}^{\ell})\right).
\]
在自回归 LLMs 中，令牌 \(x_t\) 是基于完整的历史上下文顺序生成的。我们的工作特别利用第 \(l\) 层的中间表示 \(H^{\ell}\) 来检测潜在的解码模式。

##### 推测解码。
SD 通过使用轻量级的草稿模型提议 \(k\) 个候选令牌，并由目标模型并行验证，从而加速推理。效率由接受率 \(\alpha = 1 - \mathrm{KL}(p, q)\) 和两个模型之间的成本比 \(c = T_p / T_q\) 决定，其中 \(p\) 和 \(q\) 分别代表目标分布和草稿分布。虽然像 EAGLE 或蒸馏方法这样的神经草稿通过对齐提高了 \(\alpha\)，但它们通常无论底层文本结构如何，都应用统一的草稿策略。

##### 基于上下文和自适应的推测。
除了神经草稿，基于上下文的方法如 **Prompt Lookup Decoding \(PLD\)** 使用 n-gram 匹配直接从上下文中提议候选。**CopySpec** 引入了在 n-gram 匹配和神经草稿之间切换的启发式方法。然而，这些方法仍然是与模型无关的，在“偶然重复”上会失败。最近的工作探索了统计编排：**BanditSpec** 将策略选择视为一个多臂老虎机问题，使用基于 UCB 的吞吐量反馈来优化草稿。与这些需要嘈杂探索阶段的统计“黑箱”方法不同，SwitchSD 使用“白箱”信号——探测内部状态，在推测开始前以高精度识别复制意图。

##### LLM 中的复制机制。
我们区分记忆（训练数据复制）和联想复制（上下文重现）。后者由称为**归纳头**的特定电路驱动。从概念上讲，SwitchSD 回应了 **Pointer-Generator Networks** 的门控逻辑，该网络使用学习的标量来在生成和指向之间切换。然而，我们将这个概念从生成时的词汇表门控转向推理时的推测编排。基于相关证据表明任务相关信息线性编码在 \(H^{\ell}\) 中，我们假设复制意图表现为线性可分的信号，这使我们能够训练轻量级探测器作为 SD 的高精度控制信号。

## 3 方法

### 3.1 动机：偶然重复的效率税

我们将**偶然重复**定义为当前上下文包含与过去的 n-gram 重叠，但模型的目标分布 \(p\) 与历史延续不一致的场景。如图2所示，这些情况经常源于巧合的联想模式（例如，在数学推导中重复变量名“n”），而不是结构化的复制。

在这种情况下，不加区分的复制会带来显著的**推测税**：系统在草稿提议和并行验证上花费计算预算，结果却只有 0 个令牌被接受。我们的分析显示，在推理密集型任务中，这些误报触发了超过 50% 的复制尝试。这促使我们对**潜在感知触发器**的需求——一种绕过表面巧合，直接查询模型内部表示以获取真正复制意图的方法。

### 3.2 学习复制意图探测器

借鉴复制构成一个独特解码模式的见解，我们开发了一个与任务无关的过程，直接在模型的隐藏状态中探测复制意图。与监控表面文本的启发式触发器不同，我们的探测器识别 Transformer 的底层**潜在模式**。

##### 问题表述。
我们将复制意图检测表述为一个二元分类任务。由于表面重复本质上是嘈杂的，我们对真实的训练标签采用保守的代理指标。如果一个令牌位置属于至少 \(n_{\text{train}}=5\) 个令牌的逐字序列（该序列在上下文中更早出现），则将其定义为**复制模式**实例（\(Y=1\)）。这个阈值确保我们的探测器忽略偶然的 n-gram 重叠，而专注于持续的、故意的复制行为。所有其他位置标记为非复制（\(Y=0\)）。我们的目标不是声称复制意图完全由精确的 n-gram 重叠刻画，而是使用更长的逐字跨度作为底层潜在复制模式的实用指标。关于此问题的进一步讨论见附录 E。

##### 受控探测数据集。
由于长重复 n-gram 在标准预训练语料库中很稀疏，我们构建了一个**校准刺激集**来分离复制电路。我们利用 **Claude Sonnet** 生成 1000 个提示，这些提示表现出多样化的重复结构，从结构化的代码样板到重复的语言模板。通过让目标模型为这些提示生成补全，我们提取了一个隐藏表示 \(H_j^{(\ell)}\) 及其相应复制标签 \(Y_j\) 的数据集 \(\mathcal{D}_{\ell}\)：
\[
\mathcal{D}_{\ell} := \{(H_{i,j}^{(\ell)}, Y_{i,j}) \mid i \in [1000], j \in \{M_i, \dots, N_i\}\},
\]
其中 \(M_i\) 表示解码响应的起点。这个过程提供了数十万个标记的令牌级实例，确保探测器训练具有高信噪比。附录 C 提供了数据集中具有代表性的提示，展示了它们不同程度的复制诱导行为。

##### 探测器架构和训练。
与任务相关特征线性编码在残差流中的假设一致，我们实现探测器为一个**严格线性投影**。对于每一层 \(l\)，我们训练一个独立的探测器 \(\mathcal{PR}_{\ell}\)，它直接在隐藏表示 \(H_j^{(\ell)} \in \mathbb{R}^d\) 上运行，不带偏置项：
\[
p_{\ell}(j) = \sigma(w_{\ell}^{\top} H_j^{(\ell)}),
\]
其中 \(p_{\ell}(j)\) 表示