Attention Sinks

《How Attention Sinks Keep Language Models Stable》 语言模型在处理长对话时之所以会“崩溃”，是因为当为了节省内存而移除最开始的几个token时，模型会失去稳定性。研究者发现，模型会将大量的注意力“倾倒”在最初的几个token上，把它们当作“注意力池”，用来存放那些无处安放的注意力（因为Softmax函数强制要求所有注意力权重加起来等于1）。基于此，作者提出了StreamingLLM方案，即永久保留最初的4个token，其余部分则正常使用滑动窗口，从而让模型能稳定处理超过400万个token的超长文本。这一机制已被OpenAI、HuggingFace和英伟达等公司采纳。OpenAI在其最新的开源模型中，通过为每个注意力头增加一个可学习的标量$b$来实现类似功能，其公式为：$a_i' = softmax(q_i^T K + b)$。这个简单的改动，让模型在必要时可以选择“不关注任何token”，其设计思想正是源于这项研究。 流式处理的挑战与意外发现 者在2023年暑期实习时面临一个根本性问题：如何让大模型处理超出其预训练长度的对话？直接的方案是使用“滑动窗口”，即只保留最近一部分token的键值缓存（KV Cache），丢掉最旧的。这种方法虽然高效，但实验结果却出乎意料地失败了。一旦最开始的几个token（哪怕只是一个起始符或“the”之类的普通词）被从缓存中移除，模型的困惑度（Perplexity）就会急剧飙升，输出的文本瞬间变得语无伦次。这个现象令人费解，因为这些初始token在语义上似乎并不重要，但移除它们却导致了灾难性的后果，这激发了作者深入探究其背后原因的好奇心。 “注意力池”的发现与数学原理 为了解开谜团，作者对Llama-2等模型的注意力分布进行了可视化分析，结果发现了一个惊人的固定模式：在模型的多数层中，无论当前模型在预测哪个词，总有大量的注意力会“回头看”，集中在序列最开始的几个token上。这种行为让作者联想到了图论中的“汇点”（Sink Node），即一个只进不出的节点。这些初始token就像汇点一样，吸收了来自序列中其他位置的注意力，因此作者将其命名为“注意力池”（Attention Sinks）。其背后的数学原理源于Transformer架构的核心——Softmax函数：$softmax(z)_i = \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}}$。这个函数有一个硬性约束：$\sum_{i=1}^{n} w_i = 1.0$，即所有注意力权重之和必须等于1。这意味着注意力机制不能“弃权”，必须把注意力分配到某个地方。当一个注意力头找不到特别相关的上下文时，它就会把这份“预算”投向那些因在训练中频繁出现而具有轻微偏见的初始位置，这些位置逐渐演变成专门存放闲置注意力的“泄压阀”。 问题的症结与简单的解决方案 这个发现让之前滑动窗口失败的原因豁然开朗。当移除初始token时，我们不只是丢掉了旧的上下文，而是破坏了Softmax函数分母中的一个重要部分，从而摧毁了模型维持注意力稳定的基础机制。解决方案因此变得异常简单：既然模型如此依赖这些注意力池，那我们就永远不要丢掉它们。作者提出的StreamingLLM方法，就是对传统KV缓存管理做了一个微小但关键的修改：永久保留最开始的4个“池”token，同时对其余内容使用滑动窗口。这个简单的改动带来了惊人的效果，像Llama这样原本处理几千个token就崩溃的模型，现在可以在超过400万个token的长度上保持稳定的低困惑度，实现了近乎无限长度的文本生成。 注意力池数量的探究与预训练实验 为什么恰好需要保留4个注意力池？为了回答这个问题，作者进行了一项预训练实验。他们从头训练了两个模型，一个采用标准方式，另一个则在每个训练样本的开头都加入了一个专门的[SINK]（池）token。结果揭示了一个深刻的现象：标准模型在流式处理时需要保留大约4个初始token才能稳定，因为它是在“挪用”普通内容token作为隐式的注意力池；而被投喂了专用[SINK]token的模型，在流式处理时只需要保留这1个专用的池token就足够了，而且其训练收敛效果甚至更好。这个实验证明，模型天生就需要注意力池这个功能，并且可以通过预训练让它更高效地学会使用一个专门设计的“池”，而不是临时征用多个内容token。 两种实现路径：StreamingLLM与OpenAI的设计对比 OpenAI的实现代表了解决同一问题的不同思路。作者的StreamingLLM方法是在序列开头设置一个专门的、可学习的池token $k_0$，其注意力计算方式为：$Attention(q_i, K) = softmax(q_i^T K_c, q_i^T k_0)$，其中$K_c$是内容token。而OpenAI则采用了一个更简洁的通用标量方法：$a_i' = softmax(q_i^T K + b)$`。两者的核心区别在于：StreamingLLM的设计允许不同token与“池”建立不同的关系（更具表达性），而OpenAI的设计则将“池”视为一个对所有token都一样的通用“逃生通道”（更简洁）。这体现了两种不同的设计哲学：一个是将“池”视为一个可学习的组件，另一个则是将其作为一种架构上的必需品。 注意力池的科学解释与实际应用 有研究表明，注意力池可以充当“泄压阀”，防止深度模型在处理长序列时出现“信息过度混合”（over-mixing）的病态现象，即不同token的表征变得模糊不清。注意力池通过吸引注意力，限制了信息和噪声的过度传播，从而使模型更稳健。这一效应在LLaMA 3.1 405B这样的大模型中尤为明显，其80%的注意力头都表现出注意力池现象。另一项研究则指出，如果用不要求权重和为1的其他注意力操作替换Softmax，注意力池现象就会消失，这从根本上证实了其与Softmax归一化约束的直接关系。这些理论洞见也催生了实际应用，例如CushionCache技术利用精心设计的注意力池前缀来平滑激活峰值，从而改善模型的量化效果。这也让人推测，OpenAI模型内置的注意力池机制，可能是其能够实现激进的4-bit量化的关键技术之一。 从研究到现实：一个思想的快速演进 这个始于实习期工程难题的研究，最终演变成了对Transformer架构的深刻洞见。从2023年10月英特尔集成StreamingLLM，到12月HuggingFace和苹果手机上的演示，再到2024年1月英伟达TensorRT-LLM的采纳，最后到2025年8月OpenAI开源内置注意力池的模型，这项技术在不到两年的时间里迅速从学术论文走向了全行业的主流应用。这个故事告诉我们，有时重大的科学发现并非源于宏大的理论突破，而是来自于对那些被他人忽视的、微小而奇特现象的刨根问底。通过质疑几个看似无意义的token为何如此重要，作者团队揭示了一个隐藏在所有Transformer模型中、至关重要的底层机制。