💬 NLP

Bayesian Attention Mechanism: A Probabilistic Framework for Positional Encoding and Context Length Extrapolation

本文介绍了贝叶斯注意力机制（BAM），这是一种统一现有位置编码方法的概率框架，并提出了一种广义高斯先验以实现最先进的长上下文泛化能力，从而在训练上下文长度增加 500 倍的情况下仍能实现准确的信息检索，且参数开销极小。

原作者： Arthur S. Bianchessi, Yasmin C. Aguirre, Rodrigo C. Barros, Lucas S. Kupssinskü

发布于 2026-05-11

📖 1 分钟阅读☕ 轻松阅读

原作者： Arthur S. Bianchessi, Yasmin C. Aguirre, Rodrigo C. Barros, Lucas S. Kupssinskü

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

想象一下，将 Transformer 语言模型（现代 AI 聊天机器人背后的“大脑”）比作一位试图阅读巨著的智慧超群图书管理员。这位管理员非常擅长理解词语的含义，但他们有一个奇怪的毛病：他们完全没有顺序感。如果你递给他们一个句子，比如“猫坐在垫子上”，他们无法判断“猫”是在“坐”之前还是之后，因为无论词语处于什么位置，在他们眼中看起来都是一样的。

为了解决这个问题，我们通常会赋予这位图书管理员一个位置编码（PE）。你可以将其想象成贴在每个词上的微小、无形的姓名标签或彩色贴纸，上面写着“我是第 1 个词”、“我是第 2 个词”等等。这有助于图书管理员理解故事的脉络。

然而，大多数这类“贴纸系统”都有一个缺陷：它们在处理短篇故事时表现优异，但一旦书籍变得非常长（比如一本 10 万页的小说），它们就会崩溃。图书管理员开始忽略书的开头，因为贴纸变得过于混乱或逐渐失效。

新想法：“贝叶斯注意力机制”（BAM）

这篇论文的提出者构想了一种思考这些贴纸的新方法，他们称之为BAM。BAM 不再仅仅是给词语贴上固定的标签，而是将词语的位置视为一种概率猜测。

以下是类比：
想象图书管理员试图在一份长篇文档中找到某个特定的信息（一把“通行证”）。

旧方法（如 ALiBi）： 图书管理员假设答案最有可能就在当前词语的旁边，距离越远，答案存在的可能性就越低。这就像在家里寻找丢失的钥匙；你会先检查口袋，然后是桌子，而因为车库“太远”了，你就停止在那里寻找。
BAM 方法： 图书管理员利用关于答案可能位置的“先验信念”（概率图）。这张地图不是固定的；它是一个可调整的灵活规则。

秘诀所在：“广义高斯”地图

这篇论文引入了一种特定类型的概率地图，称为广义高斯先验。你可以将这张地图想象成一种地形，其中的山丘和山谷代表了图书管理员查看某个特定词语的可能性大小。

“局部”山丘： 地图可以在当前词语旁边设有一座陡峭的山丘。这有助于图书管理员理解即时上下文（如语法和句子结构）。
“长距离”尾部： 这里是魔法所在。作者发现，通过微调一个特定的旋钮（称为 $\beta$ ），他们可以改变地图的形状。
- 如果他们将旋钮向一个方向转动，图书管理员就会忽略书中遥远的部分（就像旧方法那样）。
- 如果他们将旋钮向另一个方向转动（特别是将其设置为负数），图书管理员就会停止关注紧邻的邻居，转而 intensely 关注那些非常遥远的词语。

他们实际上取得了什么成就？

这篇论文并没有声称这能治愈疾病或为你写小说。他们在非常具体、受控的任务上进行了测试：

“大海捞针”测试： 他们将一个特定的 5 位数字（一把“通行证”）隐藏在长达数千词的文本深处。
- 结果： 当文本过长时，其他方法（如 RoPE 或 ALiBi）会放弃并随机猜测，而 BAM 模型即使面对比其训练文本长 500 倍的内容，依然能完美地找到该数字。
“困惑度”测试： 这衡量模型预测句子中下一个词的能力。
- 结果： BAM 在预测词语方面与现有最佳方法一样出色，这意味着它没有为了获得这种长距离超能力而牺牲通用的语言技能。

“隐形”的成本

他们发现的 coolest 之处在于其添加成本极低。

他们在一个已有 1.2 亿参数的模型中，仅添加了不到 1,000 个新参数（微小的记忆单元）。
这就像在沙滩上添加一粒沙子，突然整个沙滩都能看见地平线。
它既没有减慢模型的速度，也没有增加计算机资源的消耗。

用通俗英语总结

作者为 AI 构建了一套新的“位置贴纸”系统。他们不再假设答案总是在附近，而是赋予 AI 一本灵活的规则手册，使其能够忽略即时环境，并专注于深埋在过去的信息。

他们从数学上证明了这是可行的，并通过实验表明，他们的 AI 能在 50 万词的“大海”中找到“针”，而其他 AI 在几千词后就会迷失。他们并未在真实的医疗或法律文件上测试这一点，但他们证明了寻找长距离信息的机制现在变得更加强大和可靠。

技术摘要：贝叶斯注意力机制（BAM）

问题陈述

基于 Transformer 的语言模型依赖位置编码（PE）来处理 token 顺序，然而现有方法在上下文长度外推（即泛化到训练期间未见过的更长序列的能力）方面面临两个关键局限：

理论不透明性：许多 PE 技术（如正弦编码、RoPE、ALiBi、NoPE）主要基于经验动机，理论基础有限，导致其在外推期间的行为难以理解。
评估不足：过度依赖困惑度（perplexity）作为主要指标是不够的。模型可以通过局部注意力模式（仅关注最近的 token）实现低困惑度，而未能真正整合长程信息，滑动窗口评估揭示了这一缺陷。

方法论：贝叶斯注意力机制（BAM）

作者提出了贝叶斯注意力机制（BAM），这是一个理论框架，将自注意力重构为在 token 内容与位置的联合概率模型下计算值的期望。

理论框架

联合概率公式：在 BAM 中，注意力权重 $p_{ij}$ （token $j$ 关注 token $i$ 的概率）被定义为联合概率分布：
$p_{ij} = p(f_{cont}(q_i, k_j) | g_{pos}(i, j)) \times p(g_{pos}(i, j))$
其中， $f_{cont}$ 表示内容相似度， $g_{pos}$ 表示位置关系。
位置编码作为先验：该框架认为，位置编码策略对应于 token 位置上的特定先验分布 $p(g_{pos}(i, j))$ $p (g_{p os} (i, j))$ 。
- NoPE 被推导为均匀先验（受因果掩码限制）。
- ALiBi 被推导为均匀先验与拉普拉斯先验（基于距离的指数衰减）的乘积。
归一化标量（ $Z$ ）：推导引入了一个标量 $Z$ ，用于归一化内容与位置概率的乘积，代表内容与位置之间的相互作用自由能或信息耦合。

提出的解决方案：GGD-BAM

利用该框架，作者引入了GGD-BAM，它使用**广义高斯分布（GGD）**作为位置先验。添加到注意力分数上的偏置矩阵 $B$ 定义为：
$b_{ij} = -\left| \frac{j - i - \mu}{\alpha} \right|^\beta$
其中 $\mu$ 、 $\alpha$ 和 $\beta$ 是可学习参数（或针对特定行为固定）。

$\beta \in (0, 1)$ ：产生重尾分布，其衰减速度慢于拉普拉斯先验（ALiBi），允许注意力范围比 ALiBi 延伸得更远，同时保持局部性。
$\beta < 0$ ：这是对标准 GGD 定义的一种新颖松弛。这创建了一种“反转”先验，其中概率质量在查询 token 附近被抑制，并集中在上下文窗口的远端。理论上，这允许模型忽略局部上下文，专注于任意长程依赖。
可扩展 Softmax（SSMax）：作者指出，BAM 容易受到“注意力衰减”（随着上下文增长而变为均匀分布）的影响。他们将 BAM 与 SSMax 结合，后者动态重缩放 logits，以保持分布的锐度。

主要贡献

统一的理论框架：BAM 提供了现有 PE 方法的概率统一，明确将 NoPE 和 ALiBi 分别映射到均匀先验和拉普拉斯先验。
广义高斯先验：GGD-BAM 的提出引入了一类灵活的位置先验。通过学习形状参数 $\beta$ ，模型可以自适应地在局部注意力和长程检索模式之间切换。
最小开销：该方法通过为每个注意力头学习 $\theta_\alpha$ 和 $\theta_\beta$ （同时固定 $\theta_\mu = 0$ ），引入了少于 1,000 个额外的可训练参数（对于 1.2 亿参数模型，具体为 384 个）。
可解释性：该框架提供了一种原则性的方法，将注意力模式可视化和解释为位置上的显式概率分布。

实证结果

实验在 1.2 亿至 11 亿参数不等的模型上进行，这些模型在 FineWeb 10B 数据集上训练，上下文长度为 512 个 token。

密钥检索（Passkey Retrieval）：GGD-BAM（特别是结合 SSMax）在从长达32,000 个 token（训练长度的 64 倍）的序列中检索特定 token 时实现了完美准确率，并在长达512,000 个 token（训练长度的 1000 倍）的序列中保持了超过 80% 的准确率。相比之下，RoPE、正弦编码和 NoPE 等基线模型在超出其训练范围后，准确率下降至接近随机水平。ALiBi 在困惑度外推方面表现较好，但在极端长度下的检索任务中表现显著不佳。
RULER 基准（NIAH）：在“大海捞针”（Needle In A Haystack）子集中，GGD-BAM 优于所有基线（包括 RoPE 和 ALiBi），是唯一在 Single 1 任务中在 6,000 个 token 以上仍能保持高准确率的方法。
困惑度：GGD-BAM 在长上下文中的困惑度与 ALiBi 和 RoPE 相当，表明改进的检索能力并未以牺牲语言建模的流畅性为代价。
下游任务：在 MMLU 和 ARC 基准测试中，BAM SSMax 的表现优于 RoPE SSMax。
注意力可视化：可视化证实， $\beta \le 0$ 的注意力头有效地抑制了局部上下文并聚焦于远距离 token（即密钥），而 $\beta \in (0, 1)$ 的头则保持了长尾注意力。

意义与主张

该论文声称 BAM 具有双重作用：

理论层面：它为位置编码提供了严格的概率基础，阐明了为什么某些方法（如 ALiBi）在困惑度上有效，但在检索任务中却失效（由于其局部性质）。
实践层面：它提供了一种高效的方法用于上下文长度外推。通过允许注意力头学习负的 $\beta$ 值，模型可以显式地抑制局部性，从而从任意长的上下文中检索信息，而无需在训练期间直接暴露于此类长度。

作者强调，虽然困惑度是标准指标，但它不足以评估真正的长上下文能力；基于检索的任务对于验证模型是否关注远距离信息而不仅仅是依赖局部模式是必要的。这项工作表明，显式建模位置先验的能力是下一代长上下文 Transformer 的关键归纳偏置。