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

State-space language models such as Mamba and gated linear attention (GLA) offer efficient alternatives to transformers due to their linear complexity and parallel training, but often lack the expressivity and robust state-tracking needed for complex reasoning. We address these limitations by reframing sequence modelling through a probabilistic lens, using Bayesian filters as a core primitive. While classical filters such as Kalman filters provide principled state estimation and uncertainty tracking, they are typically viewed as inherently sequential. We show that reparameterising the Kalman filter in information form enables its updates to be computed via an associative scan, allowing efficient parallel training. Building on this insight, we introduce the Kalman Linear Attention (KLA) layer, a neural sequence-modelling primitive that performs time-parallel probabilistic inference while maintaining explicit belief-state uncertainty. KLA offers strictly more expressive non-linear updates and gating than GLA variants while retaining their computational advantages. On language modelling tasks, KLA matches or outperforms modern SSMs and GLAs across representative benchmarks for discrete token manipulation and state tracking.

The AI systems behind today's chatbots, called transformers, are powerful but grow slow and resource-hungry as the text they handle gets longer. Researchers have built leaner, faster alternatives, but these often struggle to keep careful track of information across long passages - a skill that genuine reasoning needs. We turned to an unlikely source: the Kalman filter, a 60-year-old algorithm that estimates the hidden state of a system from noisy measurements while keeping track of its own uncertainty. It famously helped guide the Apollo missions to the Moon and remains a gold standard in engineering and neuroscience. The catch was that it seemed to work only one step at a time - far too slowly to learn from the vast text collections modern AI needs. We found that rewriting the filter in a different mathematical form reveals hidden structure that lets all its steps be computed at once, in parallel, and built this into a neural network layer we call Kalman Linear Attention. Because the model continually reasons about its own uncertainty, useful behaviours that other systems bolt on by hand emerge on their own. It matches or improves on leading efficient models for language, and is among the first such probabilistic methods trained at large scale - a step toward AI that has a built-in sense of what it doesn't know.