🤖 machine learning

Kalman Linear Attention: Parallel Bayesian Filtering For Efficient Language Modelling and State Tracking

本文介绍了卡尔曼线性注意力（Kalman Linear Attention, KLA），这是一种可并行化的序列混合层，它通过使用信息形式的卡尔曼滤波器重新构建精确贝叶斯滤波，从而实现具有显式不确定性的非线性、扫描并行状态更新，进而超越了像 Mamba 和 GLA 这样现有的线性复杂度模型的表达能力和状态追踪能力，同时保持了计算效率。

原作者： Vaisakh Shaj, Cameron Barker, Aidan Scannell, Andras Szecsenyi, Elliot J. Crowley, Amos Storkey

发布于 2026-06-11

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Vaisakh Shaj, Cameron Barker, Aidan Scannell, Andras Szecsenyi, Elliot J. Crowley, Amos Storkey

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

核心愿景：AI“思考”的新方式

想象一下，你正在尝试阅读一本非常厚的书。

旧版 AI（Transformer）： 为了理解一个句子，AI 会查看到目前为止书中的每一个字。这就像一位图书管理员，为了找一本书，竟然把整个图书馆都搬了出来。它非常精准，但随着书的内容变长，速度会变得极其缓慢且成本高昂。
当前的效率型 AI（Mamba, GLA）： 这些模型就像一位只保留一本固定大小笔记本的图书管理员。他们只记录最重要的信息，并忘掉其余的部分。他们很快，但因为他们只是简单地将新信息“叠加”到旧信息上，所以有时会错过微妙的联系，或者在故事变得过于复杂时感到困惑。

这篇论文介绍了一位新的图书管理员：卡尔曼线性注意力（Kalman Linear Attention, KLA）。

KLA 是一位同样保留笔记本的图书管理员，但他们不仅仅是记录事实，还会记录自己对每个事实的确定程度。他们会问自己：“我记得这个，但我对此有多大的信心？这条新信息是一个改变局面的关键，还是仅仅是一个微小的细节？”

核心思想：“信心”笔记本

作者们意识到，目前的效率型 AI 模型将它们的记忆视为一个简单的数学方程（将数字相加）。但真正的思考涉及不确定性。

“信念状态”（The Belief State）： KLA 不仅仅存储一个事实，它还存储一个事实以及一个置信度分数（就像天气预报说“有 80% 的降水概率”）。
“守门员”（The Gatekeeper）： 当新信息进入时，KLA 不仅仅是将其加入。它会检查自己的信心。
- 如果 AI 对当前的记忆非常有信心，它就会忽略那些带有噪声的新信息。
- 如果 AI 不确定，它就会密切关注新信息并更新自己的记忆。
“魔法技巧”（并行化）： 通常情况下，逐步检查信心并更新记忆是很慢的（就像人们排队等待盖章一样）。这篇论文的重大突破在于证明了这种“信心检查”可以同时进行，就像一条传送带一样，这使得它能像目前最快的模型一样迅速。

创意类比

1. “怀疑派侦探” vs. “记录员”

现有模型（记录员）： 想象一位侦探，他把证人说的每一句话都写在笔记本上。如果证人说“车是红色的”，侦探就写下“红色”。如果证人后来说“其实，也许是枣红色的”，侦探就直接在旁边写下“枣红色”。笔记本变得乱七八糟，侦探可能会对到底发生了什么感到困惑。
KLA（怀疑派侦探）： 这位侦探有一个“信心计”。
- 证人： “车是红色的。”
- 侦探： “好吧，我有 90% 的把握它是红色的。我会把它记下来。”
- 证人： “等等，我觉得我也看到了一辆蓝色的车。”
- 侦察： “嗯，我对‘红色’的信心很高，但这个新线索有点不靠谱。我会稍微降低对‘红色’的信心，并记录下蓝色车的信息，但我不会立刻抹去‘红色’。”
- 结果： 侦探能够维持一个更清晰、更准确的犯罪现场图景，而不会被信息淹没。

2. “交通灯”系统

把 AI 的记忆想象成一条高速公路。

线性模型： 每辆车（新单词）都直接汇入交通流。这是一条平滑的直线。
KLA： 每辆车都有一个交通灯。
- 如果高速公路很空旷（低信心），灯就是绿灯；新车可以轻松汇入。
- 如果高速公路交通拥堵（高信心），灯就会变红；新车必须等待或非常小心地汇入。
- 创新之处： 论文研究出如何同时为一条 1000 英里的高速公路计算所有的交通灯，而不是在每一个里程碑处停下来检查灯号。

他们究竟证明了什么？

这篇论文并不声称它能治愈疾病或预测股市。它专注于语言建模（让 AI 更擅长阅读和写作）。以下是他们的演示成果：

更快且更聪明： KLA 与目前最快的模型（如 Mamba）一样快，但能解决更难的逻辑谜题。
“排列测试”： 他们给 AI 一个名为“A5”的任务，这就像一个复杂的谜题，你需要按特定顺序排列物品。
- 现有的快速模型（Mamba、标准 Transformer）除非把 AI 做得非常庞大且深层，否则无法通过这个谜题。
- KLA 用一个极小、极浅的模型就解决了它。 这证明了 KLA 比竞争对手更能追踪复杂的、变化的各种状态。
处理“长上下文”： 当 AI 需要记住 2000 字之前的剧情时，KLA 在寻找正确细节方面比其他快速模型做得更好。
规模化能力： 他们用数十亿词的数据训练了一个模型。它运行稳定，没有崩溃。这证明了这种基于“不确定性”的方法可以用于大型、现实世界的 AI 系统。

“秘诀”：OU 过程

论文提到了一个技术术语——“Ornstein-Uhlenbeck (OU) 过程”。

类比： 想象一个连接着球的橡皮筋。如果你把球拉开，橡皮筋会将它拉回中心。
在 AI 中： 这个橡皮筋防止了 AI 的“信心”失控（即爆炸到无穷大或跌落到零）。它保持了 AI 记忆的稳定性，使得模型可以进行多层堆叠而不崩溃。如果没有这个“橡皮筋”，当模型规模变大时，它会变得不稳定。

总结

卡尔曼线性注意力（Kalman Linear Attention） 是一种新型的 AI 记忆机制，它的工作方式就像一位充满信心且具有怀疑精神的侦探。它不仅仅是死记硬背，还会追踪自己对所知事物的确定程度。这使得它能比目前的快速 AI 模型更好地过滤噪声并专注于重要细节，同时运行速度同样之快。作者们证明了它在困难逻辑谜题上的有效性，并且可以在海量数据上进行训练。

技术摘要：卡尔曼线性注意力 (Kalman Linear Attention, KLA)

问题陈述

当前的深度状态空间模型（SSMs），如 Mamba 和门控线性注意力（GLA），提供了线性复杂度且可并行化的替代方案，解决了 Transformer 的二次方缩放瓶颈。然而，这些模型依赖于线性或仿射隐藏状态更新，这限制了其表达能力，无法达到由 Softmax 归一化诱导的非线性交互水平。此外，现有的线性 RNN 和 SSM 缺乏表示状态不确定性的显式机制，将隐藏状态视为确定性的。本文旨在解决这一差距：即如何在保持线性更新效率的同时，实现非线性滤波的高表达力和鲁棒性，并探讨是否可以高效地将精确的卡尔曼滤波更新实现为一种序列混合器。

方法论：卡尔曼线性注意力 (KLA)

作者提出了 Kalman Linear Attention (KLA)，这是一种将序列建模表述为精确贝叶斯滤波的概率序列混合器。其核心创新在于将卡尔曼滤波重新参数化为信息形式（使用精度和自然参数），从而实现并行计算。

1. 概率公式化

KLA 将序列混合建模为线性高斯状态空间模型中的推理过程：

动力学先验 (Dynamics Prior)： 潜在状态演化被建模为 Ornstein-Uhlenbeck (OU) 过程，并进行精确离散化。这引入了一个过程噪声项 ( $\bar{p}_t$ )，用于捕捉状态转移中的不确定性，这与标准 SSM 的确定性动力学不同。
观测模型 (Observation Model)： 输入 Token 被视为潜在状态的有噪声观测，由观测算子 ( $k_t$ ) 和代表 Token 证据置信度的值精度 ( $\Lambda^v_t$ ) 进行参数化。
读出 (Readout)： 输出通过查询算子 ( $q_t$ ) 对潜在状态的后验均值进行投影得到。

2. 通过 Möbius 变换实现并行化

经典的卡尔曼滤波本质上是顺序执行的。KLA 通过利用滤波器的信息形式克服了这一限制：

精度递归 (Precision Recursion)： 后验精度 ( $\lambda_t$ ) 的更新被证明是前一时刻精度 ( $\lambda_{t-1}$ ) 的 Möbius（分式线性）变换：
$\lambda_t = \frac{\alpha_t \lambda_{t-1} + \beta_t}{\gamma_t \lambda_{t-1} + \delta_t}$
至关重要的是，Möbius 变换具有结合律。这使得精度更新可以通过并行前缀扫描（结合律扫描）来计算，其工作量为 $O(T)$ ，深度为 $O(\log T)$ ，这与 Mamba 等现代 SSM 的并行效率相匹配。
均值递归 (Mean Recursion)： 一旦通过并行方式计算出精度路径，后验均值（信息均值）的更新也将遵循一种同样可并行的仿射变换。

3. 源自不确定性的非线性门控

在传统的线性 SSM 中，门控通常是线性或依赖于 Token 但却是确定性的；而在 KLA 中，门控自然地从精度的比例中产生。其“遗忘门”和“写入门”是依赖于历史且非线性的，由先验状态与新观测之间的相对不确定性驱动。这使得模型能够根据置信度选择性地整合信息，而无需为每个 Token 学习显式的门控参数。

4. 架构

KLA 被设计为标准 SSM 或注意力层的即插即用替代方案。它采用了类似于 Mamba 的融合 MLP 架构，结合了一维因果卷积和 SiLU 激活函数。作者实现了高效的 CUDA 和 Triton 内核，用于执行 Möbius 扫描，以避免在显存（HBM）中实例化展开后的状态。

核心贡献

结合律重参数化： 作者证明了信息形式下的对角线性高斯卡尔曼滤波具有 Möbius 结构，从而能够实现精确的并行前缀扫描进行贝叶斯滤波。
不确定性驱动的非线性： KLA 引入了源自精度比的非线性门控，在保持线性时间扫描结构的同时，提供了严格高于线性/仿射 SSM 更新的表达能力。
KLA 层： 一种实用的、即插即用的序列混合器，能够产生显式的信念状态不确定性，并可堆叠在深度网络中。
扩展与验证： 本文展示了首个在十亿级 Token 规模下训练的堆叠贝叶斯滤波原语，验证了其稳定性和性能。

实验结果

作者在合成任务、状态追踪基准测试和大规模语言建模方面对 KLA 进行了评估。

合成任务 (MAD-Lab)： 在需要处理不确定性下的选择任务（特别是 Fuzzy Recall 和 Selective Copy）时，KLA 的表现优于 GDN、GLA、M了 Mamba 和 mLSTM。它在 Compression 任务上达到了最先进的水平，表明其信息形式的滤波起到了最优上下文压缩器的作用。
状态追踪 (A5 Permutation)： 在 A5 置换组合任务（一个衡量 NC1-完全态追踪的基准测试）中，KLA 以常数深度（1–2 个 block）解决了该任务。相比之下，线性 SSM 和基于注意力的模型所需的深度随序列长度增加而增长。这表明 KLA 的分式线性更新打破了线性循环模型的表达能力极限，同时保持了可并行训练的能力。
长上下文关联回忆 (MQAR)： KLA 在序列长度为 2048 的多查询关联回忆任务中实现了近乎完美的准确率（>95%），优于 Mamba，并显著超过了在这一设置下无法学习该任务的 GLA。
大规模预训练：
- 独立运行： 在 45M 和 180M 参数规模下，独立运行的 KLA 在零样本常识基准测试（如 HellaSwag, PIQA）上达到或超过了 GPT、Mamba 和 GDN 的水平。
- 混合模式： 用 KLA 替换 GPT 的最后一个注意力层（GPT+KLA）产生的性能在混合模型中最为强劲，优于 GPT+Mamba 和 GPT+GDN。
消融实验： 作者指出 Ornstein-Uhlenbeck (OU) 先验 对于稳定的深层堆叠至关重要。如果没有这种具有均值回归特性的随机动力学，深度堆叠 KLA 会变得不稳定。

意义与主张

本文声称 KLA 成功弥合了线性状态空间模型的高效性与非线性注意力机制的高表达能力之间的鸿沟。通过将序列混合视为精确的贝叶斯滤波，KLA 提供了：

显式不确定性： 与确定性混合器不同，KLA 携带显式的信念状态不确定性，该不确定性直接控制信息门控。
更强的表达能力： Möbius 更新允许 KLA 解决线性 SSM 无法处理的状态追踪任务（如 A5），且无需牺牲并行训练能力。
可扩展性： 该研究证明了精确的贝叶斯滤波原语可以进行堆叠，并在十亿级 Token 规模下进行训练，挑战了“概率滤波在本质上是顺序的或在大型语言建模中计算成本过高”的观点。

作者将 KLA 定位为不仅是一个理论上的好奇物，而是一个实用的、即插即用的原语，它在保持长上下文处理所需的计算效率的同时，增强了现代语言模型的鲁棒性和状态追踪能力。