注意力漂移：自回归投机解码模型学到了什么

# 自回归投机解码模型学习代码的方式：https://github.com/Dogacel/Attention-Drift
来源：https://arxiv.org/html/2605.09992
作者：Doğaç Eldenk（西北大学）& Payal Mohapatra（西北大学）& Yigitcan Comlek（GE Aerospace）& Kaan Oktay（fal）& Hongyang Zhang（滑铁卢大学）& Stephen Xia（西北大学）

###### 摘要

投机解码（Speculative decoding）通过让小模型起草未来令牌来加速大语言模型（LLM）的推理。然而，起草模型（drafter models）在模板扰动和长上下文输入下性能会急剧下降。我们识别出一种此前未被报告的现象，称为**注意力漂移（attention drift）**：随着起草模型在投机链中生成连续的令牌，其注意力逐渐从提示词（prompt）转移到其自身最近生成的令牌上。这一现象在*EAGLE3*起草器和*MTP头（heads）*中均被观察到，表明漂移是起草器设计的一种固有属性。我们将此归因于链步骤之间未归一化的残差路径：随着链深度的增加，起草器的隐藏状态幅度单调增长，这表现出类似于在目标模型上堆叠额外的预归一化（pre-norm）Transformer层，而不是作为独立的自回归预测器的动态特征。为了限制这种增长，我们提出了两种架构改动：对起草器隐藏状态采用后归一化（Post-norm），以及在捕获目标隐藏状态后对每个隐藏状态应用RMSNorm。我们的干预措施使接受长度（acceptance length）相比当前领先的预归一化EAGLE3模型，在模板扰动下最高提高了2倍，在长上下文任务上提高了1.18倍，在涵盖多轮对话、数学和编程的七个标准基准上提高了1.10倍。我们的改动还使得较短的训练-测试深度能够泛化到更长的起草序列中。

![Figure 1: Attention drift.](https://example.com/image1.png)
*图1：注意力漂移。在投机过程中，起草器的注意力从提示词的“注意力汇聚点（sink token）”转移到其自身最近生成的令牌。左：在起草器的注意力热力图上展示注意力汇聚点的出现，行=查询（query），列=键（key）；颜色越深表示注意力越高。中：起草器上注意力漂移的图形化可视化。右：x轴上的每个令牌位置的注意力，虚线右侧为投机生成的令牌。*

## 1 引言

投机解码\[8, 4\]是一种无损加速大语言模型（LLM）推理的技术，其中轻量级的**起草器（drafter）**预测未来令牌，随后由较大的**目标（target）**或**验证器（verifier）**模型进行验证。在实践中，投机解码系统必须在各种部署条件下运行，包括不同的推理引擎、上下文长度、系统提示和聊天模板。在不同场景下的鲁棒性对于保持高接受率和一致的推理加速至关重要。然而，近期的研究表明，在具有挑战性的部署场景下，如模板扰动和长上下文输入，起草器的性能通常会显著下降\[17, 20\]。因此，从业者经常为特定的推理引擎、提示或模板重新训练或专门化起草器，以最大化性能。尽管投机解码模型的训练成本相对较低，但这种敏感性指出了底层存在的问题。本文的主要贡献如下：

\(i\) 我们识别出一种此前未报告的现象，有助于解释这种脆弱性，我们称之为**注意力漂移**（图1）。随着起草器在投机过程中生成更多令牌，其注意力逐渐从汇聚点（sink）转移到最近生成的令牌。我们在EAGLE-3起草器\[11\]和多令牌预测（MTP）头\[6\]中一致地观察到这一现象，表明注意力漂移反映了自回归起草器设计更广泛的属性。

\(ii\) 为了理解这种行为，我们分析了投机起草器的隐藏状态动态，发现投机步骤之间未归一化的残差连接导致隐藏状态幅度随链深度单调增长，这类似于在验证器之上堆叠额外的Transformer层。起草器隐式地学习了一个依赖于深度的细化过程，而不是一个稳定的自回归令牌预测器，使其对模板变化和长上下文敏感（第4.2节）。

\(iii\) 受此观察启发，我们引入了一种基于**后归一化（post-normalization）**（图2）的简单架构修改，并结合在隐藏状态融合之前的归一化，以防止隐藏状态幅度的增长并稳定起草过程。这些改动显著减少了注意力漂移，并将接受长度相比当前最先进的预归一化EAGLE3架构提高了最高2倍（在模板扰动下）、1.18倍（在长上下文任务上）和1.10倍（在涵盖多轮对话、推理、编程和数学的七个标准基准上）。我们进一步表明，后归一化起草器可以泛化到训练期间未见过的推理深度，从而允许较短的训练-测试深度并降低训练成本。

## 2 初步知识

投机解码交替进行两个阶段（图3）。在起草阶段，起草器自回归地生成k个候选未来令牌序列，其中k是**投机深度**。在验证阶段，目标模型在一次前向传递中评估起草的序列，并针对其自身的令牌分布应用拒绝采样；接受最长的有效前缀。一个关键的效率指标是**接受长度（τ）**，定义为每个验证轮次中接受的起草令牌的平均数量。较高的接受长度通常直接转化为固定起草器开销下的更大端到端推理加速。

EAGLE\[10\]及其后继者EAGLE-2/EAGLE-3\[9, 11\]将目标的融合内部隐藏状态通过全连接投影（fc，见图13）输入到一个单独的预归一化Transformer解码器层\[15\]中，作为起草器（图2左）。起草器使用令牌级交叉熵损失针对冻结目标预测的分布进行训练，具有固定数量的投机步骤（训练-测试深度，TTT）。EAGLE-3是我们关注的重点，是vLLM\[7\]和SGLang\[22\]等生产引擎中主导的自回归起草器设计。

**注意力汇聚点（Attention Sinks）。** 注意力汇聚点\[18\]是在推理期间吸收不成比例大量注意力的令牌（通常位于序列开头）。汇聚点在现代LLM中被广泛观察到，被认为充当稳定注意力的锚点，特别是在长上下文中。

## 3 注意力漂移

观察投机解码模型的注意力，出现了一种模式：**只要目标模型产生汇聚点，起草器就在完全相同的位置产生汇聚点。** 这一观察结果在MTP头和EAGLE-3起草器中是一致的。仔细检查具有不同注意力模式的三个模型家族后，我们观察到Llama的汇聚点在第一个令牌上，Qwen的在第二个令牌上，而GPT-oss没有汇聚点。在每种情况下，起草器都与其目标匹配。

由于起草器消耗验证器的隐藏状态，起草器的注意力类似于验证器的注意力。然而，一旦开始起草，随着起草器从消耗验证器隐藏状态转变为消耗其自身生成的隐藏状态，注意力分布逐渐发生变化。在此过程中，注意力越来越多地从汇聚点转移到最近生成的令牌。

我们将这种现象称为**注意力漂移**，在多个EAGLE-3起草器中可视化于图4和图5。在具有强汇聚点的模型中，汇聚点在起草过程中逐渐减弱，而注意力质量集中在最近起草的令牌上。因此，起草器在两种模式下运行：在预填充（pre-fill）时关注验证器的隐藏状态，在生成时关注其自身的隐藏状态，并在两者之间进行渐进的学习过渡。我们假设这种过渡使得起草器**对分布外输入更加脆弱**，最明显的是在长上下文和更深的投机步骤中。

这里，我们检查MTP头\[6\]，作为跨架构检查，以确定注意力漂移是EAGLE-3特有的还是消耗其先前输出的起草器的更一般属性。MTP头是辅助预测头，在预训练期间与目标模型联合训练，共享目标的LM头并预测多个位置之后的令牌。这与EAGLE-3不同，后者是在冻结目标之上作为单独的起草器进行事后训练，并具有自己的LM头。

我们检查了Qwen3.5 9B的MTP头（图6），这是一个单层Transformer，遵循目标模型的注意力架构，并在连续的投机步骤中重用相同的权重。在MTP投机过程中，我们观察到类似的趋势（图7）：对汇聚点令牌的注意力显著减少，而对最近起草令牌的注意力增加。

**如果验证器没有注意力汇聚点怎么办？**
一些近期模型，如Qwen3-Next和GPT-oss，旨在抑制注意力汇聚点。Qwen3-Next使用门控注意力，对SDPA输出应用每头sigmoid门，以便头可以乘法抑制其贡献。GPT-oss则在softmax分母中引入每头可学习的偏差logit，为每个头提供一个明确的“不关注任何内容”选项，以吸收多余的注意力质量。检查GPT-oss起草器（图8），我们观察到第一个令牌上没有可见的汇聚点，但在重复的模板令牌上存在微弱但一致的注意力集中。这表明，汇聚点类似的行为不仅可能源于架构归纳偏置，还可能源于重复的特殊令牌或模板标记，特别是在验证器缺乏专用架构汇聚点机制的情况下。

## 4 什么导致了注意力漂移？

为了理解为什么发生注意力漂移以及它揭示了关于底层模型的什么信息，我们从三个不同的验证器模型家族训练了起草器。我们主要在MT-Bench\[21\]上进行评估，八个类别的多轮指令使我们能够在多样化的提示上进行评估。本节中探索的所有架构更改都需要重新训练起草器（详细信息见附录C）。

### 4.1 隐藏状态幅度

注意力漂移发生在我们从目标的隐藏状态向生成过程中消耗起草器自身的隐藏状态转移时。因此，我们首先检查验证器和起草器的隐藏状态。在表1中，比较了验证器的隐藏状态、验证器的融合隐藏状态以及不同投机步骤中起草器自身隐藏状态的幅度，使用均方根（RMS），$\|x\|_2/\sqrt{d}$。我们注意到接下来描述的几个模式。

**表1：跨越模型家族的隐藏状态幅度（RMS），平均取自80个MT-Bench提示和k=8个投机轮次。** 来自验证器低、中、高层的隐藏状态（$h_{low}, h_{mid}, h_{high}$）；验证器隐藏状态的融合表示$h_{FC}$；以及投机步骤1、2、3和8中的起草器输出隐藏状态（$h^*_1$到$h^*_8$）

| 验证器隐藏状态 | | | | 起草器输出 $h^*_t$ | | | |
| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |
| **目标模型** | $h_{low}$ | $h_{mid}$ | $h_{high}$ | $h_{FC}$ | k=1 | k=2 | k=3 | k=8 |
| Llama 3.1 8B | 0.56 | 0.58 | 0.78 | 12.46 | 3.92 | 4.87 | 5.86 | 14.02 |
| Qwen 3 30B | 0.33 | 2.17 | 2.71 | 2.56 | 1.00 | 1.13 | 1.25 | 1.67 |
| Qwen 3.5 35B | 0.03 | 0.12 | 0.21 | 0.89 | 3.47 | 4.03 | 4.42 | 5.92 |
| GPT-oss 20B | 35 | 23 | 24 | 96 | 87 | 89 | 95 | 138 |
| GPT-oss 120B | 316 | 314 | 555 | 834 | 975 | 1153 | 764 | 76 |

**观察1：隐藏状态幅度不匹配。**
验证器、目标和全连接层的隐藏状态幅度差异很大。这归因于起草器中RMSNorm的位置以及$h_{FC}$和起草器$h^*$之间共享的归一化分量（如图13所示）。尽管目标和起草器的隐藏状态幅度差异巨大，但我们没有看到隐藏状态幅度与模型接受率之间的相关性，这提示我们模型学会了对此进行补偿。

**观察2：验证器融合表示$h_{FC}$不平衡。**
我们探索的所有三个目标家族都使用预归一化架构：RMSNorm应用于每个Transformer块内部（在注意力和MLP之前），但从未应用于残差流本身。因此，每一层的幅度通过深度单调累积（预归一化稀释\[14\]），对于表1的每一行都满足$\|h_1\| < \|h_2\| < \|h_3\|$。目标的最终预LM头RMSNorm通常在生成logits时吸收这种增长，但EAGLE-3使用在此归一化之前捕获的验证器内部状态（$h_{low}, h_{mid}, h_{high}$）。因此，$h_{high}$以其更大的幅度主导了$h_{FC}$信号。因此在未来的章节中，我们在每个用作生成$h_{FC}$输入的目标隐藏流之前添加RMSNorm层，以防止这种不平衡并生成更稳定的验证器表示。

**观察3：跨越投机深度的幅度增长。**
起草器生成的隐藏状态随投机深度单调增长。在所有模型家族中，$h^*_k$远离分布累积幅度