aiqbench analysis
Kalman Linear Attention (KLA) resolves the sequential bottleneck of classical Kalman filters by reparameterizing updates into an associative scan, enabling temporally parallel probabilistic inference for language modeling. This architecture strictly outperforms existing linear state-space models like Mamba and GLA in expressivity and state tracking while maintaining linear computational complexity. As a drop-in primitive trained at scale, KLA offers a significant advancement for applications requiring robust uncertainty estimation and complex permutation tasks.
Kalman Linear Attention: Parallel Bayesian Filtering For Efficient Language Modelling and State Tracking
arXiv:2602.10743v2 Announce Type: replace Abstract: 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\"obius/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.
Abstract: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.
Accepted at ICML 2026. An earlier version of this work was presented at the 1st Workshop on Epistemic Intelligence in Machine Learning (EIML) at EurIPS 2025
Machine Learning (cs.LG)
arXiv:2602.10743 [cs.LG]
(or arXiv:2602.10743v2 [cs.LG] for this version)
Full text under the publisher's open licence · 225 words
Related coverage
Kalman Delta Networks: Uncertainty-aware Associative Memory
arXiv · Sep 9, 2026
PIRAMID: Progress and Plans
LessWrong · Jul 27, 2026
Introducing Instella-MoE: A State-of-the-Art Fully Open Mixture-of-Experts Language Model
AMD · Jul 24, 2026
Serve Kimi-K2.5-MXFP4 on MI355X with ATOM
AMD · Jul 23, 2026