🤖 machine learning

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

This paper introduces Kalman Linear Attention (KLA), a parallelizable sequence mixing layer that reformulates exact Bayesian filtering using an information-form Kalman filter to achieve non-linear, scan-parallel state updates with explicit uncertainty, thereby surpassing the expressivity and state-tracking capabilities of existing linear-complexity models like Mamba and GLA while maintaining computational efficiency.

Original authors: Vaisakh Shaj, Cameron Barker, Aidan Scannell, Andras Szecsenyi, Elliot J. Crowley, Amos Storkey

Published 2026-06-11

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Vaisakh Shaj, Cameron Barker, Aidan Scannell, Andras Szecsenyi, Elliot J. Crowley, Amos Storkey

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: A New Way for AI to "Think"

Imagine you are trying to read a very long book.

Old AI (Transformers): To understand a sentence, the AI looks at every single word in the book so far to decide what the next word should be. It's like a librarian who pulls out the entire library to find one book. It's very accurate but gets incredibly slow and expensive as the book gets longer.
Current Efficient AI (Mamba, GLA): These are like a librarian who keeps a small, fixed-size notebook. They only write down the most important stuff and forget the rest. They are fast, but because they just "add" new notes to the old ones, they sometimes miss subtle connections or get confused if the story gets too complex.

This paper introduces a new librarian: Kalman Linear Attention (KLA).

KLA is a librarian who keeps a notebook, but instead of just writing facts, they keep track of how sure they are about every fact. They ask themselves: "I remember this, but how confident am I? Is this new piece of information a game-changer, or just a small detail?"

The Core Idea: The "Confidence" Notebook

The authors realized that current efficient AI models treat their memory like a simple math equation (adding numbers together). But real thinking involves uncertainty.

The "Belief State": KLA doesn't just store a fact; it stores a fact and a confidence score (like a weather forecast saying "80% chance of rain").
The "Gatekeeper": When new information comes in, KLA doesn't just add it. It checks its confidence.
- If the AI is very confident in its current memory, it ignores the new, noisy information.
- If the AI is unsure, it pays close attention to the new information and updates its memory.
The Magic Trick (Parallelism): Usually, checking confidence and updating memory step-by-step is slow (like a line of people waiting to get a stamp). The paper's big breakthrough is showing that this "confidence check" can be done all at once, like a conveyor belt, making it just as fast as the current speediest models.

Creative Analogies

1. The "Skeptical Detective" vs. The "Note-Taker"

Current Models (The Note-Taker): Imagine a detective who writes down everything a witness says in a notebook. If the witness says "The car was red," the detective writes "Red." If they later say "Actually, maybe it was maroon," the detective just writes "Maroon" next to it. The notebook gets messy, and the detective might get confused about what actually happened.
KLA (The Skeptical Detective): This detective has a "confidence meter."
- Witness: "The car was red."
- Detective: "Okay, I'm 90% sure it's red. I'll write that down."
- Witness: "Wait, I think I saw a blue car too."
- Detective: "Hmm, my confidence in 'red' is high, but this new clue is shaky. I'll lower my confidence in 'red' slightly and note the blue car, but I won't erase the red car yet."
- Result: The detective maintains a much clearer, more accurate picture of the crime scene without getting overwhelmed.

2. The "Traffic Light" System

Think of the AI's memory as a highway.

Linear Models: Every car (new word) just merges into the traffic flow. It's a smooth, straight line.
KLA: Every car has a traffic light.
- If the highway is clear (low confidence), the light is green; the new car merges easily.
- If the highway is jammed with heavy traffic (high confidence), the light turns red; the new car has to wait or merge very carefully.
- The Innovation: The paper figured out how to calculate all these traffic lights for a 1,000-mile highway simultaneously, rather than stopping at every single mile marker to check the light.

What Did They Actually Prove?

The paper doesn't claim this will cure diseases or predict the stock market. It focuses on language modeling (making AI smarter at reading and writing). Here is what they demonstrated:

It's Faster and Smarter: KLA is just as fast as the current speediest models (like Mamba) but can solve harder logic puzzles.
The "Permutation" Test: They gave the AI a task called "A5," which is like a complex puzzle where you have to shuffle items in a specific order.
- Current fast models (Mamba, standard Transformers) failed this puzzle unless the AI was made huge and deep.
- KLA solved it with a tiny, shallow model. This proves KLA can track complex, changing states better than its competitors.
It Handles "Long Contexts": When the AI has to remember a story from 2,000 words ago, KLA did a better job at finding the right details than the other fast models.
It Works at Scale: They trained a model with billions of words of data. It didn't crash. It worked stably, proving that this "uncertainty-based" approach can be used for big, real-world AI systems.

The "Secret Sauce": The OU Process

The paper mentions a technical term called the "Ornstein-Uhlenbeck (OU) process."

Analogy: Imagine a rubber band attached to a ball. If you pull the ball away, the rubber band pulls it back toward the center.
In the AI: This rubber band prevents the AI's "confidence" from going crazy (exploding to infinity or crashing to zero). It keeps the AI's memory stable, allowing it to stack many layers deep without breaking. Without this "rubber band," the model would become unstable when made larger.

Summary

Kalman Linear Attention is a new type of AI memory that works like a confident, skeptical detective. It doesn't just memorize; it tracks how sure it is about what it knows. This allows it to filter out noise and focus on important details much better than current fast AI models, all while running just as quickly. The authors proved it works on hard logic puzzles and can be trained on massive amounts of data.

Technical Summary: Kalman Linear Attention (KLA)

Problem Statement

Current state-space models (SSMs) like Mamba and gated linear attention (GLA) offer linear-complexity, parallelizable alternatives to softmax attention, addressing the quadratic scaling bottleneck of Transformers. However, these models rely on linear or affine hidden state updates, which limits their expressivity compared to the nonlinear interactions induced by softmax normalization. Furthermore, existing linear RNNs and SSMs lack explicit mechanisms for representing state uncertainty, treating hidden states as deterministic. This paper addresses the gap between the efficiency of linear updates and the expressivity/robustness of nonlinear filtering, asking whether exact Kalman filter updates can be implemented efficiently as a sequence mixer.

Methodology: Kalman Linear Attention (KLA)

The authors propose Kalman Linear Attention (KLA), a probabilistic sequence mixer that formulates sequence modeling as exact Bayesian filtering. The core innovation lies in reparameterizing the Kalman filter into information form (using precisions and natural parameters) to enable parallel computation.

1. Probabilistic Formulation

KLA models sequence mixing as inference in a linear-Gaussian state-space model:

Dynamics Prior: The latent state evolution is modeled as an Ornstein-Uhlenbeck (OU) process, discretized exactly. This introduces a process noise term ( $\bar{p}_t$ ) that captures uncertainty in state transitions, distinct from the deterministic dynamics of standard SSMs.
Observation Model: Input tokens are treated as noisy observations of the latent state, parameterized by an observation operator ( $k_t$ ) and a value precision ( $\Lambda^v_t$ ) representing confidence in the token evidence.
Readout: The output is derived from the posterior mean of the latent state, projected via a query operator ( $q_t$ ).

2. Parallelization via Möbius Transforms

Classical Kalman filtering is inherently sequential. KLA overcomes this by exploiting the information form of the filter:

Precision Recursion: The update of the posterior precision ( $\lambda_t$ ) is shown to be a Möbius (fractional-linear) transformation of the previous precision ( $\lambda_{t-1}$ ).
$\lambda_t = \frac{\alpha_t \lambda_{t-1} + \beta_t}{\gamma_t \lambda_{t-1} + \delta_t}$
Crucially, Möbius transformations compose associatively. This allows the precision updates to be computed via parallel prefix scans (associative scans) with $O(T)$ work and $O(\log T)$ depth, matching the parallel efficiency of modern SSMs like Mamba.
Mean Recursion: Once the precision path is computed in parallel, the posterior mean (information mean) updates follow an affine transformation that is also parallelizable.

3. Nonlinear Gating from Uncertainty

Unlike linear SSMs where gating is often linear or token-dependent but deterministic, KLA's gating emerges naturally from the ratio of precisions. The "forget gate" and "write gate" are history-dependent and nonlinear, driven by the relative uncertainty between the prior state and the new observation. This allows the model to selectively incorporate information based on confidence without requiring explicit, learned gating parameters for every token.

4. Architecture

KLA is designed as a drop-in replacement for standard SSM or attention layers. It follows a fused MLP architecture similar to Mamba, incorporating 1D causal convolutions and SiLU activations. The authors implement efficient CUDA and Triton kernels for the Möbius scan to avoid materializing expanded states in high-bandwidth memory (HBM).

Key Contributions

Associative Reparameterization: The authors demonstrate that diagonal linear-Gaussian Kalman filtering in information form admits a Möbius structure, enabling exact parallel prefix scans for Bayesian filtering.
Uncertainty-Driven Nonlinearity: KLA introduces nonlinear gating derived from precision ratios, offering strictly greater expressivity than linear/affine SSM updates while maintaining linear-time scan structure.
KLA Layer: A practical, drop-in sequence mixer that produces explicit belief-state uncertainty and can be stacked in deep networks.
Scaling and Validation: The paper presents the first stacked Bayesian-filtering primitives trained at the billion-token scale, validating their stability and performance.

Experimental Results

The authors evaluate KLA across synthetic tasks, state-tracking benchmarks, and large-scale language modeling.

Synthetic Tasks (MAD-Lab): KLA outperforms GDN, GLA, Mamba, and mLSTM on tasks requiring selection under uncertainty, specifically Fuzzy Recall and Selective Copy. It achieves state-of-the-art results on Compression, suggesting its information-form filtering acts as an optimal context compressor.
State Tracking (A5 Permutation): In the A5 permutation-composition task (a benchmark for NC1-complete state tracking), KLA solves the task with constant depth (1–2 blocks). In contrast, linear SSMs and attention-based models require depth that grows with sequence length. This demonstrates that KLA's fractional-linear updates break the expressivity limits of linear recurrent models while remaining parallel-trainable.
Long-Context Associative Recall (MQAR): KLA achieves near-perfect accuracy (>95%) on the Multi-Query Associative Recall task with sequence length 2048, outperforming Mamba and significantly exceeding GLA, which fails to learn the task in this setting.
Large-Scale Pretraining:
- Standalone: At 45M and 180M parameters, standalone KLA matches or improves upon GPT, Mamba, and GDN on zero-shot commonsense benchmarks (e.g., HellaSwag, PIQA).
- Hybrid: Replacing the final attention layer of a GPT with KLA (GPT+KLA) yields the strongest performance among hybrids, outperforming GPT+Mamba and GPT+GDN.
Ablation: The authors show that the Ornstein-Uhlenbeck (OU) prior is critical for stable deep stacking. Without the mean-reverting stochastic dynamics, training deep stacks of KLA becomes unstable.

Significance and Claims

The paper claims that KLA successfully bridges the gap between the efficiency of linear state-space models and the expressivity of nonlinear attention mechanisms. By casting sequence mixing as exact Bayesian filtering, KLA provides:

Explicit Uncertainty: Unlike deterministic mixers, KLA carries an explicit belief-state uncertainty that directly controls information gating.
Strictly Higher Expressivity: The Möbius updates allow KLA to solve state-tracking tasks (like A5) that linear SSMs and attention cannot, without sacrificing parallel training capabilities.
Scalability: The work demonstrates that exact Bayesian filtering primitives can be stacked and trained at the billion-token scale, challenging the notion that probabilistic filtering is inherently sequential or computationally prohibitive for large-scale language modeling.

The authors position KLA not merely as a theoretical curiosity but as a practical, drop-in primitive that enhances the robustness and state-tracking capabilities of modern language models while maintaining the computational efficiency required for long-context processing.