← 返回 2026-03-13

注意力 sink 在 softmax Transformer 中是理论上必需的:来自触发条件任务的证据 Attention Sinks Are Provably Necessary in Softmax Transformers: Evidence from Trigger-Conditional Tasks

Yuval Ran-Milo 📅 2026-03-12 👍 2 2026-07-13 08:36
ReLU注意力 Transformer架构 softmax归一化 注意力机制 理论分析

证明 softmax 注意力的归一化约束使 sink 成为触发条件任务的必要机制

前置知识

Attention Sink

注意力 sink 指在 Transformer 中,注意力概率质量集中在一个固定的、与内容无关的位置上的现象。通常这个位置是序列开头的 BOS(beginning of sequence)token。例如,在 Llama 2-7B 等模型中,某些注意力头在没有特定触发时会将近乎全部注意力分配给第一个位置,这种行为被称为 sink。这种现象在各种规模模型和不同位置编码方案(如 ALiBi、RoPE)下都被观察到。

理解注意力 sink 是阅读本文的核心基础。本文的核心结论就是论证这种现象不是训练的副产品,而是 softmax 注意力的数学必然性。如果不理解什么是 sink 以及它在真实模型中的表现,就无法理解作者要解决什么问题以及证明的意义。

Softmax Attention

Softmax 注意力是 Transformer 中最核心的计算组件。给定查询向量 Q、键向量 K 和值向量 V,注意力权重计算为 $\alpha_{i,j} = \frac{\exp(Q_i K_j^\top)}{\sum_{k=1}^i \exp(Q_i K_k^\top)}$,其中分母是所有前面位置的指数分数之和,确保了权重的归一化(概率分布)。这种归一化使得每个位置对前面的所有 token 分配的注意力权重和为 1,这是与 ReLU 等非归一化注意力的关键区别。

本文的核心论点就是 softmax 的归一化约束是导致 sink 形成的根本原因。理解 softmax 的数学性质(特别是概率分布约束)对于理解作者的证明过程至关重要,因为整个理论分析都围绕归一化如何迫使模型在需要产生默认输出时必须找一个稳定锚点展开。

Trigger-Conditional Behavior

触发条件行为指模型在检测到特定触发信号时执行复杂计算,而在其他情况下执行无操作。真实模型中的例子包括:Gemma 7B 中的'撇号头'在遇到撇号时会聚合上下文,否则注意力塌缩到 BOS;Llama 2-7B 中的活跃-休眠头在代码输入时展现多样化注意力模式,在文本输入时塌缩到 sink。这种行为可以建模为函数 f(x) 在触发条件 T(x) 成立时执行复杂操作,否则输出默认值。

本文的理论分析完全建立在一个精心设计的触发条件任务上,这个任务模拟了真实模型中观察到的触发条件行为。理解这个概念是理解作者为什么要设计这样一个合成任务、以及这个任务如何反映真实模型运作机制的前提。

ReLU Attention

ReLU 注意力是 softmax 注意力的一种替代方案,用 ReLU 激活函数替代了 softmax 归一化。具体定义为 $\alpha_{i,j} = \text{ReLU}(Q_i K_j^\top) / n_i$,其中 $n_i = \max(i-1, 1)$ 是归一化因子。与 softmax 不同,ReLU 注意力的权重不需要和为 1,这意味着模型可以通过将所有注意力权重设为零来直接输出零向量,而不需要一个 sink 来吸收概率质量。本文证明了 ReLU 注意力可以在没有任何 sink 的情况下解决相同的触发条件任务。

ReLU 注意力是本文用来证明 sink 不是任务结构必然性的关键对比工具。作者通过展示 ReLU 注意力可以零 sink 解决相同任务,强有力地证明了 softmax 的归一化约束(而非任务本身)是 sink 形成的根本原因。理解 ReLU 注意力的运作机制对于理解这个对比论证至关重要。

研究动机

Transformer 模型普遍存在注意力 sink 现象,即注意力概率质量集中在一个固定的、与内容无关的位置(通常是序列开头的 BOS token)。这种现象在小到大规模模型中都被广泛报告,并且出现在各种位置编码方案下,包括绝对/学习嵌入、ALiBi、RoPE 甚至没有显式位置编码的情况。更严重的是,sink 会带来实际危害:当概率质量集中在固定位置时,注意力可能从其他 token 被转移开,影响下游准确率;在压缩和量化时会加剧数值问题;扭曲基于注意力的可解释性分析;使流式和长上下文推理复杂化。尽管有这么多负面影响,sink 存在的根本原因尚不清楚——是训练过程的副产品,还是在某些情况下功能上必需的?

本文的目标是本文的核心目标是从理论角度明确回答注意力 sink 的存在原因。具体来说,作者希望证明在某些设置下,sink 不是优化或训练 regimes 的意外产物,而是 softmax Transformer 为了实现特定功能(触发条件行为)而必须具备的结构特征。为了实现这个目标,作者设计了一个理论分析框架,包括一个精心构造的触发条件任务、两种注意力机制(softmax 和 ReLU)的对比分析,以及一系列必要性定理和构造性定理。最终目标是揭示 softmax 归一化约束如何迫使注意力塌缩到稳定锚点,并找到避免 sink 的可能路径。

与已有工作不同的是,本文的独特切入角度是从理论必要性而非经验相关性角度研究 attention sink。此前的工作如 Barbero et al. (2025) 认为 sink 通过 BOS 作为稳定锚点实现了受控信息混合;Guo et al. (2024) 分析了 sink 形成的训练动态。但这些工作都是经验性的,没有证明 sink 在某些情况下是不可避免的。本文的 gap 在于:通过严格的数学证明,建立 sink 在 softmax 注意力中的理论必要性,同时通过 ReLU 注意力的构造性反例,证明 sink 不是任务结构的必然结果。这种对比分析首次从理论层面揭示了归一化约束作为 sink 形成根本原因的角色,为理解注意力机制的本质提供了全新的视角。

核心方法

本文的方法框架包含三个核心组成部分:理论分析、构造性反例和实验验证。理论分析部分设计了一个触发条件任务,该任务要求模型在检测到触发 token 时输出前面所有 token 表示的平均值,在其他位置输出零向量。这个任务模拟了真实模型中观察到的'有触发时聚合上下文,无触发时休眠'的行为模式。然后作者在这个任务上建立了 softmax 自注意力的两个必要性定理:对于单层模型,任何接近最优的解必须在所有非触发位置将近乎全部注意力分配给 BOS token;对于多层模型,至少有一层必须在某个非触发位置展示 sink 行为。构造性反例部分通过显式构造一个 ReLU 注意力模型,证明相同的任务可以在没有任何 sink 的情况下完美解决。实验验证部分在单层和多层、单头和多头的设置下训练模型,可视化注意力模式,验证理论预测。

核心创新点是从概率单纯形上的归一化必须迫使注意力塌缩到稳定锚点这一直觉出发,建立严格的数学必要性证明。与已有研究本质区别在于:本文不是观察性地描述 sink 现象,而是证明在特定任务下,sink 是数学上不可避免的。作者通过设计一个最小化的触发条件任务,剥离了所有无关复杂性,使得归一化约束的效应可以被清晰分析。另一个关键创新是通过与 ReLU 注意力的对比,证明驱动 sink 形成的是 softmax 的归一化约束而非任务结构本身。这种对比分析有力地分离了不同因素的影响,为理解注意力机制的本质提供了清晰的因果证据。

方法步骤详情

方法的完整步骤分为四个阶段。第一阶段:任务设计。作者构造输入分布 D,序列长度为 L,每个 token 表示包含四个坐标类型:(i) BOS 指示器,仅对第一个 token 为 1;(ii) 触发指示器,仅对触发 token 为 1;(iii) 非触发非 BOS 指示器,对所有其他 token 为 1;(iv) 内容坐标,从连续分布独立采样。触发位置 j 从 {2, ..., L} 均匀采样。目标输出 y(i) 在所有位置为零向量,除了触发位置 i=j,其中 y(j) = (j-1)^{-1}\sum_{k=2}^j x(k)(前面非 BOS token 的平均值)。损失函数为最坏情况损失 L(f) = sup_{(x,y) \in support(D)} max_{i \in [L]} \|y(i) - f(x)(i)\|。第二阶段:单层必要性证明(定理 1)。假设存在一个 softmax 注意力模型 f 在此任务上实现 L(f) \leq \eta,用反证法证明当 \eta \to 0 时,对于任何非触发位置 i \neq j,必须有 \alpha_{i,1} \geq 1 - \epsilon。证明思路是:如果 \alpha_{i,1} \leq 1 - \epsilon 在某个非触发位置以正概率成立,那么一个常数量的注意力质量落在非 BOS token 上;通过鸽巢原理和独立采样,这会导致值映射必须压碎一个正概率集的 token,在触发位置产生矛盾。第三阶段:多层必要性证明(定理 2)。将多层网络展开,应用类似推理:如果没有层展示 sink 行为,内容 token 上的有效注意力权重仍然很大,迫使值映射将它们压碎为零,与触发位置所需的敏感性矛盾。第四阶段:ReLU 构造(定理 3)。显式构造一个 ReLU 注意力模型,通过选择查询和键权重与触发指示器坐标和非触发非 BOS 指示器坐标对齐,使得注意力权重在触发位置为正常数(计算平均值),在其他位置为零。由于 ReLU 不强制归一化,模型可以通过零注意力权重直接输出零向量,不需要 sink。

技术新颖性

技术新颖性体现在三个方面:第一,首次建立了 attention sink 的理论必要性证明。此前的工作都是经验观察,而本文通过严格的数学分析,证明了在 softmax 注意力中,为了解决触发条件任务,sink 是数学上不可避免的。这是对注意力机制理论的重要贡献。第二,通过对比 softmax 和 ReLU 注意力,首次将归一化约束识别为 sink 形成的根本原因。这种因果分离分析在注意力机制研究中是创新性的,为理解不同注意力变体的行为差异提供了理论基础。第三,提出的触发条件任务框架是一个理论工具,可用于研究更广泛的注意力现象。这个任务虽然合成,但抓住了真实模型中观察到的关键行为模式('触发时活跃,否则休眠'),为理论分析和实验验证提供了一个清晰的测试平台。整体而言,本文的新颖性在于将注意力机制研究从描述性提升到了因果性和必然性层面。

实验结果

实验结果强有力地验证了理论预测。在单层单头模型的实验中(定理 1 和定理 3 验证),作者可视化了 1000 个测试样本上的平均注意力权重,触发位置固定在位置 8。对于 softmax 注意力,模型在每个非触发位置将近乎单位质量的注意力分配给位置 1(BOS token),方差可以忽略不计(图 3a 和 3b),完全符合定理 1 的预测。对于 ReLU 注意力,模型在位置 1 上的注意力权重在整个序列中保持接近零(图 3c 和 3d),没有形成 sink,完全符合定理 3 的预测。在多层多头模型的实验中(定理 2 验证),2 层 2 头的 softmax 模型的所有头都在非触发位置展示强 sink 行为(图 4)。在更深的模型中,作者观察到与定理 2 的存在性预测一致的模式:sink 在某些但不是所有头中出现。例如,在一个达到低损失的 4 层 4 头 softmax 模型中,第 4 层第 3 个头在所有非 BOS 位置上对 BOS 的注意力接近零,而同一网络中的其他头发展出清晰的 sink(图 6 和附录图 7)。这确认了定理 2 的存在性本质:网络中某处必须存在 sink,但不是在每个头中。替换 softmax 为 ReLU 注意力完全消除了多层模型中的 sink 形成:图 5 显示 2 层 2 头 ReLU 模型的所有头都没有发展出 sink,4 层 4 头模型也是如此(附录图 8)。这些实验结果共同验证了理论预测,并证明这些模式超越了理论分析的设置:softmax 模型发展出强 sink 而 ReLU 注意力在单层和多层变体中都消除了它们。

定理 1 和定理 3 的验证
Figure 3: 定理 1 和定理 3 的验证
多层多头验证。2 层 2 头 softmax 模型在随机输入上的注意力模式(触发在位置 8)
Figure 4: 多层多头验证。2 层 2 头 softmax 模型在随机输入上的注意力模式(触发在位置 8)
ReLU 注意力:2 层 2 头模型
Figure 5: ReLU 注意力:2 层 2 头模型
多层 Transformer 中没有 sink 的 softmax 头
Figure 6: 多层 Transformer 中没有 sink 的 softmax 头
查看结构化数据
任务指标本文基线提升
触发条件平均值计算(合成任务) 注意力 sink 强度(非触发位置对 BOS 的注意力权重) Softmax 单层:接近 1.0;ReLU 单层:接近 0.0;Softmax 多层:至少一个头接近 1.0;ReLU 多层:所有头接近 0.0 无基线(这是理论必要性证明,不是性能比较) 不适用(本文证明的是数学必要性,不是性能提升)
触发条件平均值计算(合成任务) 任务准确率(最坏情况损失) Softmax 和 ReLU 模型都能达到低损失(接近最优) 无基线 不适用
注意力模式稳定性(方差测量) 注意力权重标准差 Softmax:可忽略方差(稳定 sink);ReLU:一致行为 无基线 不适用

局限与改进

作者承认的局限性包括:首先,合成的触发条件任务虽然经验上基于真实 sink 行为,但仍然是一个简化设置。实际任务可能更复杂, sink 的必要性可能在不同设置下有所不同。其次,对于多层模型,必要性结果(定理 2)只保证至少有一层必须在某个非触发位置展示 sink 行为,但没有刻画具体哪一层必须如此。作者的经验观察显示 sink 确实不是在所有位置和层中形成,但理解 sink 确切在哪里出现需要动态分析优化如何在有效解中选择,这留给未来工作。第三,本文的分析适用于广泛的注意力机制类别(满足归一化和单调性),但可能有其他架构(如门控注意力或 Mamba-based 模型)不需要 sink。最后,本文主要关注数学必要性和存在性证明,没有深入研究 sink 在实际应用中的具体影响或缓解策略的有效性。我自己的观察是:本文的任务设置可能过于理想化(内容坐标从连续分布独立采样),真实数据的复杂结构和相关性可能影响 sink 的必要性;此外,论文没有考虑不同初始化、优化器、学习率等超参数对 sink 形成的影响,虽然理论证明与优化无关,但实际训练中可能有额外因素。

独立分析的弱点

独立分析的第一个弱点是理论结果的普适性有限。论文证明的 sink 必要性仅适用于特定的触发条件任务,虽然作者声称分析可以几乎直接扩展到更广泛的触发条件问题(如键查询检索),但没有提供严格的泛化证明。改进方向是建立更广泛的任务分类,明确哪些任务类别需要 sink,哪些不需要。第二个弱点是对于实际应用的指导意义有限。论文证明了 sink 在理论上不可避免,但没有提供在不破坏任务性能的前提下缓解 sink 副作用的具体方案。改进方向是结合理论约束设计新的注意力机制,既能满足归一化要求又能避免 sink 的负面效应,例如引入可学习的 'off' 状态表示或混合归一化方案。第三个弱点是实验验证的范围有限。虽然实验验证了理论预测,但只在合成任务上进行,没有在真实任务(如语言建模、机器翻译)上验证理论洞察的适用性。改进方向是在更广泛的基准任务上测试不同注意力机制的 sink 行为和性能关系,提供更全面的经验证据。

未来方向

作者提出的未来工作方向包括:首先,形式化刻画需要 sink 的完整任务类别。本文的触发条件任务及其变体(如键查询检索)代表了更广泛的触发条件问题中的一个特定计算模式。作者认为分析很可能扩展到相关任务,但将完整类别的形式化描述留给未来工作。其次,研究 sink 在其他架构中的行为,包括其他特殊稳定且总是存在于输入中的 token(如推理模型中的 <|think|>)是否展示类似的 sink 行为,以及相对新发现的'二级注意力 sinks'(secondary attention sinks, Wong et al., 2026)现象。第三,理解优化如何在多层模型的有效解中选择,特别是 sink 确切在哪里出现。本文的定理保证存在性,但刻画具体位置需要动态分析优化过程。基于本文成果可延伸的研究方向包括:设计sink 免费注意力机制,直接支持无操作操作;研究 sink 与模型规模、深度、宽度的关系;探索 sink 在不同模态(视觉、多模态)中的异同;开发基于理论约束的 sink 缓解策略,确保不破坏触发条件电路的功能。

复现评估

论文的复现性评估如下:代码方面,作者提供了代码仓库:https://github.com/YuvMilo/sinks-are-provably-necessary,这是有利的。数据方面,论文使用的是合成任务,数据通过算法生成,不需要外部数据集,这大大简化了复现过程。实验设置方面,论文提到使用序列长度 L = 16,训练细节在附录 A 中,但没有提供完整的超参数列表(如学习率、优化器、训练轮数等)。算力方面,任务相对简单,应该不需要大规模计算资源。难度方面,复现理论验证实验应该相对直接,但复现完整的证明过程需要深入理解数学分析。总体而言,复现难度中等,主要挑战在于从有限的训练细节描述中推断完整的实验设置,以及正确实现合成任务和不同的注意力机制。代码仓库的存在大大降低了复现门槛。