用卡尔曼滤波实现并行化语言建模,提升状态追踪能力。
Kalman Linear Attention: Parallel Bayesian Filtering For Efficient Language Modelling and State Tracking
- 将卡尔曼滤波重参数化为信息形式,实现并行更新
- 在相同计算量下比GLA更强大,可解决排列组合任务
- 适合需要精确状态估计的长序列建模场景
状态空间语言模型如Mamba和门控线性注意力(GLA)提供线性复杂度、可并行化的替代方案,但其线性状态更新限制了表达能力和鲁棒性。本文从概率视角出发,将序列混合视为精确贝叶斯滤波,以卡尔曼滤波为核心构建。传统卡尔曼滤波被视为串行过程,我们通过将其重参数化为信息形式,使更新变为关联扫描,从而实现每标记的非线性(摩比乌斯/精度递归)但时间并行的更新。由此产生的卡尔曼线性注意力(KLA)层是即插即用的序列混合器,支持时间并行的概率推理,显式携带信念状态不确定性,且在同等计算成本下比GLA更富表达力。该表达力直接转化为更强的状态追踪:KLA解决了线性状态空间模型与注意力无法处理的排列组合(A5)任务,同时保持扫描并行。作为即插即用原语,在合成标记操作和零样本常识基准上表现匹配或优于现代状态空间模型和GLA,并成为首批在十亿级标记规模训练的堆叠贝叶斯滤波原语之一。
原文摘要 · Abstract (English)
State-space language models such as Mamba and gated linear attention (GLA) offer linear-complexity, parallelisable alternatives to transformers, but their linear state updates limit expressivity and robust state tracking. We close this gap from a probabilistic angle, casting sequence mixing as exact Bayesian filtering with the Kalman filter as the core primitive. Classical Kalman filters give principled state and uncertainty estimates but are viewed as inherently sequential; we show that reparameterising them in information form turns their updates into an associative scan - so the per-token recurrent update is non-linear (a Möbius/precision recursion) yet remains temporally parallel. The resulting Kalman Linear Attention (KLA) layer is a drop-in sequence mixer that performs time-parallel probabilistic inference, carries an explicit belief-state uncertainty, and is strictly more expressive than GLA-style linear updates at the same computational cost. This expressivity translates directly into stronger state tracking: KLA solves permutation-composition ($A_5$) tasks that linear SSMs and attention cannot, while staying scan-parallel. As a drop-in primitive it also matches or improves on modern SSMs and GLAs across synthetic token-manipulation and zero-shot commonsense benchmarks, and is among the first stacked Bayesian-filtering primitives trained at the billion-token scale.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。