Manish Gupta 解析 Attention Sinks 现象及其在超长文本语言建模中的应用

Attention Sinks 的现象与成因

Manish Gupta 在视频中首先通过 Llama-2（70亿参数）模型的注意力热力图，揭示了自回归模型中一个反直觉的现象。在序列长度为 16 的输入下，观察第 0 层和第 1 层的注意力分布时，模型对初始 token 的关注度较低，且模式更偏向局部。然而，随着层数加深（如第 2、9、16、23 乃至第 31层），第一个 token（即起始标记）的注意力值变得极高，呈现为热力图中的“全红”状态。这种现象表明，在深层网络中，初始 token 成为了所有后续 token 的注意力汇聚点。Gupta 指出，这些被过度关注的初始 token 被称为 Attention Sinks（注意力沉陷）。其根本原因在于 Softmax 函数的归一化约束：当某个词在语义上自足、无需关注其他 token 时，Softmax 仍强制将注意力质量分配给某些位置。由于自回归建模的特性，这些未被利用的注意力质量最终“沉陷”到了最常被观察到的初始 token 上，导致模型被迫在每一层都保留对起始标记的高强度关注。

不同注意力机制的性能对比

针对长文本处理，传统的密集注意力（Dense Attention）因计算复杂度随序列长度平方增长而效率低下。例如，若输入长度为 1024，则需进行约 $10^6$ 次计算。为了解决效率问题，窗口注意力（Window Attention）仅缓存并关注前 L 个 token。然而，Gupta 强调 滑动窗口注意力（Sliding Window Attention）与常规窗口注意力的关键区别在于是否进行重新计算。在滑动窗口模式下，模型不依赖缓存的 Key-Value 对，而是基于当前滑动的窗口重新计算注意力。实验数据显示，当窗口足够大时，带重计算的滑动窗口注意力在准确率上显著优于普通窗口注意力。这是因为普通窗口注意力的困惑度（Perplexity）在初始 token 被移出缓存时会急剧上升，而滑动窗口通过重计算维持了上下文的一致性。Gupta 指出：“有趣的是，如果窗口足够长，带重计算的滑动窗口实际上会导致比窗口注意力显著更低的困惑度。”

Streaming LLM 与超长序列建模

基于 Attention Sinks 的发现，Streaming LLM 提出了一种高效的注意力机制。既然初始 token 至关重要且必然承载大量注意力质量，Streaming LLM 不再仅关注局部窗口，而是显式地计算对前几个 token（即 Attention Sinks）的注意力。这种方法将时间复杂度从序列长度的平方级降低为 输入长度乘以窗口大小 的线性级别。在针对高达 20,000 token 序列长度的测试中，Streaming LLM 在保持较低困惑度（即高准确率）的同时，显著降低了延迟。相比于密集注意力的高计算开销和窗口注意力的精度损失，Streaming LLM 通过利用 Attention Sinks 的特性，实现了效率与精度的平衡。Gupta 总结道：“Streaming LLM 是一个简单的想法：不仅要关注局部上下文，还要关注前几个 token。”这一机制使得语言模型能够有效地处理 400 万以上的超长输入序列，突破了传统预训练长度对模型能力的限制。

传统滑动窗口在超长序列下的性能崩溃

当输入序列长度超过预训练长度时，传统的模型表现会出现显著退化。以 Llama 2-7B 为例，Streaming LLM 的困惑度（Perplexity）值保持平稳且较低，即使序列长度达到 20k。相比之下，“dension”（推测为 Dense Attention）模型在输入序列长度超过预训练长度后，困惑度会突然变得非常高。更糟糕的是，使用滑动窗口注意力（Sliding Window Attention）的模型在输入序列长度大于 KV 缓存大小时，前几个初始令牌会被从缓存中驱逐（evicted），导致性能急剧下降。由于 Llama 2 使用 RoPE 嵌入，这种效应甚至更为强烈。

"the dension guy basically suddenly has a very high perlexity as soon as uh your input sequence length goes beyond the pre-training length."

Streaming LLM 的核心机制：保留注意力同步令牌缓存

Streaming LLM 解决“注意力同步”（Attention Sinks）问题的核心策略是：保留前四个注意力同步令牌（attention sync tokens）的键值缓存，并结合滑动窗口缓存。例如，在生成第 9 个令牌时，模型不仅关注上下文中的 6、7、8，还始终关注最初的四个令牌（0, 1, 2, 3）。为了实现这一点，必须重新索引位置 ID：不再将缓存中的令牌标记为 0,1,2,3,6,7,8，而是重新映射为连续的 0-5。这使得模型感觉这些同步令牌并非遥远的前缀，而是直接附加在当前滚动 KV 缓存之前的部分。

"keep the key value uh caches of the four attention sync tokens... you basically pay attention to them always in that senses."

推理效率与同步令牌数量的实证分析

在保持计算复杂度与滑动窗口近似的同时，Streaming LLM 比带有重计算的滑动窗口方法快 22 倍。实验显示，在 Llama、Pythia、Falcon 和 MPT 等不同模型及不同规模下，即使序列长度达到百万级别（4M tokens），困惑度依然保持低位。关于同步令牌的数量，研究通过表格数据对比了使用 0、1、2、4 个同步令牌的效果：

此外，滑动窗口的大小对最终效果影响不大。即使将 Llama 的窗口大小从 4,000 减少到 2,000，甚至进一步降至更小的值，困惑度并未显著升高。这表明注意力同步机制在小窗口下依然有效。

预训练阶段的改进：引入全局可训练同步令牌

除了推理时的调整，研究还探索了在预训练阶段指定第一个令牌为全局可训练的注意力同步嵌入（global trainable designated attention sync token）。通过反向传播随机初始化并训练该令牌，观察注意力分布的变化：

• 无同步令牌：在低层（如 Layer 0, 2），注意力主要集中在局部；随着层级加深，高层（如 Layer 10+）对初始令牌的注意力显著增加。
• 有同步令牌：从 Layer 0 开始，注意力就高度集中在该同步令牌上。随着层级加深，对其他初始令牌的注意力逐渐减少。

这种设计类似于视觉 Transformer 中使用的 Register Tokens，旨在通过预训练引导模型更有效地利用全局上下文信息。

引入Attention Sink机制优化稀疏注意力

传统的流式语言模型（Streaming LLM）若采用朴素的全局或局部注意力，往往面临困惑度（Perplexity）居高不下的问题。引入可学习的同步Token（Sync Token）后，模型不再需要缓存前几个Token，而是可以通过稀疏注意力或基于窗口的注意力机制仅依赖一个同步Token来维持性能。实验表明，使用同步Token后，困惑度值从第一个Token开始就显著下降。相比之下，如果不使用该机制，模型的 perplexity 会非常高。这种基于 Attention Sink 的技术不仅有效降低了长序列建模的难度，还不会导致下游任务精度的任何下降。在多种自然语言处理任务的评估中，引入同步Token的模型表现稳定，没有任何任务出现性能受损的情况。这证明了该机制在保持长序列建模能力的同时，兼顾了下游应用的准确性。

Streaming QA基准测试中的长上下文能力验证

在流式问答（Streaming QA）基准测试中，研究者将 ARC Challenge 和 ARC Easy 数据集的所有问答对拼接成单个超长序列进行预测，并评估精确匹配准确率（Exact Match Accuracy）。对于 Llama 2 的 7B、13B 和 70B 参数模型，密集注意力（Dense Attention）在如此长的序列中直接导致内存溢出异常。而采用窗口注意力的模型虽然能运行，但准确率表现较差。相反，流式语言模型（Streaming LLM）在这些测试中表现出极高的准确性。值得注意的是，这些准确率与单样本提示（One-shot）下的表现相当，即模型在单独接收每个示例时能达到的高水平准确率。这表明流式评估保留了长序列中的关键信息提取能力，即使将所有示例拼接为单一输入也能维持高性能。这一结果有力地证明了Attention Sinks机制在处理超长上下文时的有效性。

Stream Eval基准测试与窗口注意力的局限性分析

Stream Eval 基准测试进一步验证了模型在长序列中的记忆能力。该任务要求模型记住20行之前的内容，并在后续每增加10行的查询中持续回忆对应位置的信息。测试结果显示，流式语言模型在长序列下保持了极高的准确率。然而，密集注意力模型的准确性在其输入长度超过预训练窗口时突然下降；而窗口注意力的准确率则在输入长度超出 KV Cache 大小、导致早期Token被驱逐时开始下降。正如视频所述："> window attention worsens when the initial tokens get excluded... as soon as your input length goes beyond the KV cache size"。虽然窗口注意力结合重计算（Recomputation）可以保留准确率，但这种方法非常昂贵且耗时。相比之下，流式语言模型通过添加注意力同步Token与最近Token结合的方式，解耦了预训练窗口大小与实际文本生成长度。即使处理400万Token的序列，该方法依然表现优异。虽然预训练时使用多个同步Token（如4个）可能比单个略好，但核心思想——利用少量注意力同步Token配合局部窗口注意力来建模超长序列——已被证明是极具价值的创新。

结论

视频论证表明，利用注意力沉陷（Attention Sinks）机制的 Streaming LLM 能有效解决超长序列建模中的精度与效率矛盾。通过保留前四个同步令牌的 KV 缓存，该模型在序列长达百万级时仍保持低困惑度及高问答准确率，且推理速度比带重计算的滑动窗口快 22 倍。相比之下，传统密集注意力面临内存溢出风险，而普通滑动窗口在初始令牌被驱逐后性能急剧下降。尽管引入全局可训练同步令牌进一步优化了注意力分布，但该机制主要依赖特定缓存策略而非完全稀疏化。总体而言，该方法在兼顾长上下文记忆能力与计算效率方面表现优异，未对下游任务精度造成损害。