💬 NLP

You Need Better Attention Priors

本文介绍了 GOAT，这是一种通过利用熵正则化最优传输（Entropic Optimal Transport）将标准注意力机制中的隐式均匀先验替换为可学习的连续先验，从而实现对标准注意力的泛化，进而解决注意力汇聚（attention sinks）问题并实现具备长度泛化能力的空间编码，同时保持与 FlashAttention 等优化内核的兼容性，这是一种新颖的注意力机制。

原作者： Elon Litman, Gabe Guo

发布于 2026-08-25

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Elon Litman, Gabe Guo

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

在现代人工智能的广阔版图中，一种被称为 Transformer 的特定类型计算机程序已成为一些现存最强大的语言和图像系统的引擎。在这些系统的核心，存在着一种被称为“自注意力”（self-attention）的机制，它允许软件决定在任何给定时刻应该关注句子或图像的哪些部分。想象一位正在阅读长篇文档的读者；自注意力就是那种让读者能够瞬间将代词如“它”关联回其所指代的特定名词，或者将动词与其主语联系起来的心理过程，无论两者之间间隔了多少个单词。多年来，工程师们一直依赖一套标准且略显僵化的规则来引导这一聚焦过程。这些规则在许多情况下表现良好，但当文本变得非常长，或者系统遇到从未见过的模式时，它们就会显得力不从心，往往导致混乱或对关键信息的关注丧失。

斯坦福大学的一个研究小组提出了一种关于这种聚焦机制运作方式的根本性转变。他们认为，标准的规则基于一个隐藏的假设，即计算机在查看内容之前，应该认为每一个可能的单词或图像块都有同等的可能重要。研究人员指出，这是一个过于简化的起点。相反，他们开发了一种名为 GOAT 的新方法，允许系统学习属于自己的“预期”去何处寻找。GOAT 并非强迫计算机从一张白纸开始，而是让它发现并采用特定的结构化习惯，例如一种自然倾向于关注句子开头或最近单词的习惯，而无需受到单词实际含义的干扰。

这一发现的核心在于重新构想注意力——它不仅仅是一个简单的相似度计算，而是一个在序列项中分配注意力的过程。在标准方法中，系统会尝试尽可能均匀地分散注意力，除非内容强烈暗示了其他方向。研究人员发现，这种“均匀分布”的假设正是系统在内容模糊或序列极长时失效的原因。通过用一个灵活的、可学习的引导取代这种均匀假设，新方法使系统即使在信息稀疏的情况下也能保持稳定的焦点。这对于被称为“注意力汇聚”（attention sinks）的现象尤向尤为重要：在长对话中，模型往往会将不成比例的注意力倾倒在少数几个特定的标记（tokens）上，通常是第一个标记，仅仅是因为系统没有其他地方可以放置其注意力。新方法将这种行为解释为一种逻辑结果，即当系统缺乏明确信号时如何分配注意力，并提供了一种有意识控制这种行为的方法。

为了测试他们的想法，研究人员构建了一个将注意力的结构规则与单词的实际含义分离的系统。在以往的方法中，寻找目标的规则往往与单词的含义纠缠在一起，使得系统难以推广到新的长度或语境。新方法将这两者区分开来。它学习一个连续的预期映射，该映射可以在系统训练过程中进行调整。这个映射可以捕捉复杂的模式，例如对紧邻当前单词的前一个单词的偏好，或者对序列开头的偏向，而不会扭曲单词本身的含义。研究人员证明，该系统可以使用驱动当前大型语言模型的高速计算机芯片进行高效实现，不需要额外的内存或处理能力。

实验结果令人瞩目。当他们在海量文本上训练模型时，新方法在理解长序列方面表现优于现有技术。在要求模型在长上下文中寻找特定隐藏信息的测试中（这项任务通常被称为“大海捞针”），即使上下文长度是训练期间所见长度的许多倍，新系统仍能保持近乎完美的准确度。相比之下，其他依赖固定位置规则处理方法的性能随着文本增长而迅速下降。新系统在模拟生物序列（如 DNA）以及处理图像方面也表现出了卓越的性能，在图像处理中，它学会了处理二维空间关系，而无需被明确告知如何操作。

其中一个最重要的发现是系统如何处理“注意力汇聚”问题。在标准模型中，过度关注第一个标记的倾向通常是模型试图理解数据的副作用，这有时会干扰其处理新信息的能力。研究人员表明，通过明确教导系统将第一个标记作为一种结构规则的默认焦点，实际上可以提高稳定性。这使得模型可以将第一个标记作为一个可靠的锚点，当其余文本不清晰时使用，而不会破坏其他单词的表示。这种结构与含义的分离意味着系统能够比以前更优雅地适应新的长度和语境。

研究人员还探索了当输入数据发生意外变化时，这种新方法是如何表现的。在一个实验中，他们在一个短序列上训练模型，然后将其测试在长十六倍的序列上。虽然其他方法性能显著退化，但新系统仍能保持高准确度运行。这表明，该系统已经学习了一套稳健的组织信息规则，可以应用于远超其训练数据的场景。这种高效泛化的能力是构建能够处理现实世界人类交流中开放式、变长特性的人工智能的关键一步。

在计算机视觉领域，该方法同样展现了多功能性。当应用于图像识别任务时，系统学会了理解图像不同部分之间的空间关系。它发展出一种观察相邻图像块的偏好，这种模式模仿了人类扫描视觉场景的自然方式。这使得系统即使在从未见过的分辨率下也能识别物体和场景，而这种能力在其他模型中通常难以实现，除非进行广泛的重新训练。研究人员发现，系统自发地发现了这些空间偏差，证实了该方法不仅限于文本，也可以应用于任何具有序列或空间结构的数据。

这项工作还为为什么某些行为会在大型语言模型中出现提供了更清晰的理论解释。通过将注意力框架化为一个平衡内容与学习到的预期的过程，研究人员解释了为什么模型有时在长上下文中表现挣扎，以及如何解决它。他们表明，这些系统中常见的稳定性问题并不是架构本身固有的缺陷，而是使用固定且无信息起始点的后果。通过允许系统学习自己的起始点，他们创造了一个更稳定、更可靠的基础。

这项研究并不声称已经解决了人工智能的所有问题，但它为构建更稳健的系统提供了一个强大的新工具。该方法旨在直接替换当前最先进模型中使用的注意力机制，这意味着它可以以极小的改动集成到现有软件中。作者已公开其代码，允许其他研究人员测试并基于其发现进行开发。随着领域向能够处理更长上下文和更复杂任务的模型迈进，控制和理解注意力如何分布的能力可能会成为一项标准要求。新方法提供了一种比以往更具灵活性和精确性的方式来实现这一点，为开发不仅更聪明而且更稳定、更可靠的推理系统铺平了道路。

这项工作的意义超越了仅仅提升模型的阅读或视觉能力。它为机器如何组织信息提供了一个全新的视角。通过将注意力的规则视为可以被学习和完善的东西，而非由工程师硬编码的规则，研究人员为构建能够根据特定任务需求调整其内部结构的系统打开了大门。这种从僵化规则向灵活学习的转变是人工智能发展中的重要一步，让我们离能够以人类般的轻松应对现实世界复杂性的系统又近了一步。研究结果表明，构建更强大机器的关键可能不在于使其规模更大，而在于赋予它们更好的聚焦方式。

技术摘要：你需要更好的注意力先验

问题陈述

标准的自注意力机制虽然在经验上取得了成功，但其依赖的启发式方法缺乏统一的数学基础。具体而言，缩放点积注意力（scaled dot-product attention）的动机是将点积作为一种相似度度量，并将 softmax 作为平滑的 argmax 替代方案。这种公式化形式隐含地假设了对 Token 位置的均匀先验（uniform prior），除非内容得分有正当理由，否则会对偏离均匀分布的行为进行惩罚。这种假设导致了三个主要局限性：

表达能力与稳定性之间的权衡： 试图提高表达能力的各种方法（例如相对位置编码）往往会将结构偏差与语义内容表示纠缠在一起，从而导致不稳定性。
长度泛化： 标准注意力在处理显著长于训练时见的序列长度时表现挣扎。
注意力汇聚（Attention Sinks）： 在长上下文场景中，模型表现出“注意力汇聚”现象，即不成比例地将概率质量分配给特定的 Token（通常是第一个 Token）。虽然这通常被视为一种学习到的人工痕迹，但本文认为这是当前公式化的一种结构性必然，当被迫通过内容向量来学习时，会破坏语义表示。

方法论：熵正则最优传输 (EOT) 与 GOAT

理论基础：注意力即 EOT

作者通过熵正则最优传输 (EOT) 的视角重新审视了自注意力。在这种视角下，注意力权重是传输问题的一个解：即如何分布一个单位质量（查询 Query）到各个键（Keys）上，以最小化传输成本（负亲和力）并最大化熵。

标准注意力： 对应于由香农熵（Shannon entropy）正则化的 EOT 问题，这在数学上等同于最小化与均匀先验之间的 Kullback-Leibler (KL) 散度。
泛化： 作者提出用一个任意的可学习先验分布 $\pi$ 取代隐含的均匀先验。这使得目标转变为最小化与 $\pi$ 之间的 KL 散度。

理论结果： 对于先验 $\pi$ ，最优传输方案是基于内容得分并经过对数先验偏移后的 softmax：
$p^*_j = \text{softmax}\left(\frac{s_j}{\tau} + \log \pi_j\right)$
其中 $s_j$ 是未缩放的点积得分， $\tau$ 是温度参数。这表明标准的相对位置编码（PEs）仅仅是这种由 EOT 导出的先验的启发式近似。

GOAT 机制

为了高效地实现这一理论，作者引入了 GOAT（Generalized Optimal transport Attention with Trainable priors，具有可训练先验的广义最优传输注意力）。GOAT 将对数先验 $K_{ij}$ 参数化为一个关于 Token 位置的连续、可微函数，并满足三个标准：平移不变性（相对性）、全局默认值（汇聚点）以及与优化核（如 FlashAttention）的兼容性。

核心组件：

谱相对先验（Spectral Relative Prior）： 相对分量 $K^{rel}_{ij}$ $K_{ij}^{r e l}$ 使用截断傅里叶级数进行参数化。这使得模型能够学习在特定频率下的对称（吸引）和反对称（排斥）相互作用。
- 它通过三角恒等式分解为查询（Query）和键（Key）向量，使其可以作为标准的点积进行计算，而无需显式生成 $L \times L$ 的偏差矩阵。
- 这确保了平移等变性，从而实现了对未见序列长度的外推。
显式汇聚参数化（Explicit Sink Parameterization）： 为了解决注意力汇聚问题，GOAT 引入了一个专门的仅针对键的对数偏置（key-only logit bias） $u(j)$ $u (j)$ 。
- 它被实现为键向量上的广播偏置，且独立于查询（Query）。
- 这允许模型建立一个结构性默认值（例如，关注第一个 Token），而不会通过膨胀内容向量的范数来破坏语义表示，从而实现解耦。
统一实现： GOAT 构建了复合的查询和键向量，将内容子空间与位置子空间进行拼接。通过精心缩放这些向量，加性的对数先验被注入到标准的缩放点积注意力核中，而无需修改内核本身或增加计算开销。这保持了与 FlashAttention 及其他 I/O 感知内核的完全兼容性。

核心贡献

数学泛化： 本文将注意力形式化为一个 KL 正则化的传输问题，证明了标准注意力是均匀先验的一个特例。它推导出了带有任意先验注意力的闭式解。
GOAT 架构： 一种新的注意力机制，能够学习连续的先验，结合了学习型嵌入的灵活性与固定编码的长度外推鲁棒性。
对注意力汇聚的理论解释： 作者从第一性原理出发解释了注意力汇聚，将其解释为低信号机制下 EOT 目标的最佳解。他们证明了可以通过先验显式建模汇聚点，从而将其从语义内容中解耦，以提高稳定性。
效率与兼容性： GOAT 被实现为标准多头注意力的掉换式替代方案，无需修改优化后的内核，并保持线性时间复杂度。

实验结果

作者在语言建模、长上下文检索、生物序列建模和计算机视觉任务中验证了 GOAT。

语言建模与外推：
- 在 C4 数据集（125M 参数）上，GOAT 在分布内困惑度（perplexity）方面优于 ALiBi（提升 1.55 点），同时在 16 倍训练长度的外推方面保持稳健，而 RoPE 则会出现灾难性的性能下降。
- 在使用 FineWeb 训练的 1.23B 参数模型中，GOAT 匹配了 RoPE 的分布内性能，但在外推方面显著优于后者（例如在 4 倍和 8 倍训练长度下），而 RoPE 的困惑度会剧烈飙升。
长上下文检索：
- 在 Passkey 检索和针尖寻 haystack (NIAH) 任务中，GOAT 在远超训练长度的上下文长度下仍能保持近乎完美的准确率，而基准模型（RoPE、ALiBi、正弦 PE）则表现出显著退化。
生物序列建模：
- 在人类参考基因组序列上，与 RoPE 相比，GOAT 实现了更低的验证负对数似然（NLL）和更好的生成质量（通过 GC% 轨迹对齐衡量）。
- 它还比 RoPE 减少了 36% 的峰值 CUDA 内存分配。
计算机视觉：
- 应用于 ImageNet-1k 上的 Vision Transformer，GOAT 学习到了一个二维平移不变先验，实现了对更高输入分辨率（最高达 512px）的鲁棒零样本外推，而标准绝对位置嵌入在此类场景下会失效。

意义与主张

本文声称，标准自注意力的脆弱性（表现为较差的长度泛化能力和涌现的注意力汇聚）源于其隐含的均匀先验假设。通过用可学习且具有表达力的先验取代这一假设，GOAT 实现了三个核心进步：

解决了表达能力与稳定性的权衡： 它结合了学习型先验的适应性与静态启发式方法的鲁棒外推能力。
对汇聚点的理论辩护： 它将注意力汇聚重新定义为一种最优传输默认值，而非需要修复的 Bug，可以对其进行显式建模和控制，而不必破坏语义表示。
实际效率： 它在标准的优化注意力内核内实现了这些理论收益，作为一种掉换式的改进方案，为现代 Transformer 提供了极小的计算开销。

作者将 GOAT 不仅仅定位为一种新的位置编码，而是作为注意力机制本身的泛化。这表明，如果正则项足够具有表达力，那么通过操纵内容向量来编码几何结构的需求就会变得多余。