# Kalman State Space Models, Simply Explained

#100 May 4 Theory

https://www.aibutsimple.com/p/kalman-state-space-models-simply-explained

Hello from the AI, but simple team! If you enjoy our content and custom visuals, consider sharing this newsletter with others or [upgrading](http://www.aibutsimple.com/pro?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=kalman-state-space-models-simply-explained) so we can keep doing what we do.

For decades, the Kalman filter has been widely used for estimation theory, providing a mathematically elegant method to track hidden states through noisy observations.

At the same time, deep sequence models have evolved from simple [recurrent networks](https://newsletter.aibutsimple.com/p/recurrent-neural-networks-explained-mathematically) to [transformer architectures](https://newsletter.aibutsimple.com/p/transformers-and-the-attention-mechanism) and, most recently, [selective state space models](https://newsletter.aibutsimple.com/p/mamba-and-state-space-models-simply-explained).

The Kalman area and the sequence model path have largely developed side-by-side, sometimes borrowing from each other but never fully combining.

The [Kalman-Optimal Selective State Space Model (KOSS)](https://arxiv.org/pdf/2512.16723?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=kalman-state-space-models-simply-explained) (Wang et al., 2025) changes that.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/24b9c6df-57c0-4fd1-a953-4ebe826f23e9/Continuous-time_SSM_KOSS__2_.png)

KOSS offers a strong mathematical combination of these two areas.

- By using Kalman filters to estimate what important details to remember in a sequence (the selection problem), it’s able to be grounded in mathematics rather than heuristic design choices.

The result is a model that is not only theoretically motivated but practically better.

It consistently outperforms state-of-the-art baselines across nine forecasting benchmarks while using fewer parameters and less memory than its competitors.

### What Are State Space Models?

When a model reads a sequence (a sentence, a time series, or a stream of sensor data), it should do two things simultaneously: remember what happened earlier and decide what's worth remembering.

A State Space Model (SSM) is one way to do this. Think of it like a conveyor belt with a memory box.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/1f0804f4-3e4a-464d-9c0b-02a9bb472d80/State_Space_Model__SSM__Block_Diagram__2_.png)

As each new token or data point arrives, the model updates its memory box based on what just came in, then uses that memory to make predictions or pass information forward.

The "state" is just a compact summary of everything the model has seen so far, like a running tally. The "space" refers to the mathematical framework that governs how that summary evolves over time.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/a7c25362-8f0b-4290-9388-825ce7746aaa/State_Space_Model__SSM__Block_Diagram__1_.png)

SSMs became popular in machine learning because they are much faster than Transformers on long sequences.

Transformers compare every token to every other token (which gets expensive fast), while SSMs compress the past into a fixed-size state and update it step by step, like reading a book and keeping a rolling summary rather than re-reading every page from the start.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/2cd1068e-2c77-4085-9130-6c5de24181a0/Copy_of_What_is_a_State_Space_Model.png)

The latest generation of SSMs, like Mamba, added a crucial twist: selective updating. Instead of updating the memory the same way regardless of what comes in, the model learns to assign more or less importance to each new input.

- In the above equation, this means that the “B” matrix is determined at every forward pass, dynamically changing based on the input.

This made SSMs significantly more powerful. But as KOSS reveals, Mamba's selection mechanism has a fundamental blind spot, and fixing it is what KOSS is all about.

---

### [Level Up Your Career With AI, But Simple PRO](https://www.aibutsimple.com/pro?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=kalman-state-space-models-simply-explained)

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/281697a8-c915-435c-a68b-43374b671698/aibutsimple_academy.png)

Supporting us by purchasing a PRO membership will get you:

- Exclusive Hands-on Tutorials With Code
- Real-World Interview Questions and Cases (From Real Companies)
- Full Math Explanations and Beautiful Visualizations
- Ad-Free Reading for an Uninterrupted Learning Experience

[Break Into AI With PRO](https://www.aibutsimple.com/pro?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=kalman-state-space-models-simply-explained)

---

### What KOSS Fixes

To understand what KOSS contributes, it helps to understand first what it fixes.

The key innovation in Mamba was input-dependent selection: rather than using fixed transition matrices, the model computes its parameters (A, B, C) as functions of the current input.

- This allows the model to selectively retain or discard information based on what it is currently seeing, allowing it to solve tasks like selective copying.

But input-dependent selection has a structural blind spot. When the model computes its selection weights, it only looks at the current input token.

It has no mechanism for asking: “Given what I have seen so far, given the context in my latent state, does this input look relevant or suspicious?”

This matters enormously in real-world sequences, where correlated noise and distractors mimic genuine signals at the token level but are distinguishable from them over time.

- For instance, maybe an individual token seems relevant, but the model has no way to notice that it’s duplicated with 2 other tokens next to it.

This is the fundamental problem KOSS is designed to solve: selection should be governed not only by the current input but also by the history compressed into the latent state.

We can observe the efficacy with a sample context-aware task. The researchers used a sequence containing tokens and correlated distractors, which are similar to the real tokens.

- Mamba's accuracy collapses from near-perfect to around 13.5% at 50% interference.
- KOSS maintains 79.2% accuracy on the same task because its selection mechanism can distinguish genuine signals from plausible distractors by consulting the accumulated context in its latent state.

### Kalman Filters and The Math Behind Them

The central contribution of KOSS is a reframing of the selection problem.

KOSS looks at what context is currently compressed within the latent state and how it should update that context given a new observation to minimize uncertainty. It doesn’t select input importance based on only the token content itself.

This is done by using a Kalman filter. In classical Kalman filtering, the system maintains a prior estimate ĥ-(t) of the hidden state, observes a new measurement x(t), and forms a posterior distribution estimate (what it believes) through:

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/90280a64-04cf-42e7-a6fb-ca9aee8160bd/Your_paragraph_text__1_.png)

The term (x(t) - Cĥ-(t)) is called the innovation (which is central to KOSS), and it is the difference between the observed measurement and what the model predicted it would be.

- It is exactly zero if the model's prediction is perfect and is large when the observation is surprising.

The Kalman gain K(t) then determines how much to trust this innovation. A high gain means the update is dominated by the new observation, and a low gain means it is dominated by the prior belief.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/8b0114b1-12c1-4431-9d19-19e93a8fce33/Kalman_Gain__KOSS_.png)

The Kalman gain is not chosen arbitrarily. It is the optimal solution that minimizes the expected squared estimation error, shown below.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/20d617c4-b5c6-49a2-b2c3-e313a61f0eb0/image.png)

Above, e(t) = h(t) - ĥ(t) is the estimation error. Minimizing the trace of the resulting error covariance yields the closed-form solution:

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/c9733186-b97f-4e3e-802b-c9c38923fd14/image.png)

Here, P-(t) is the prior error covariance and R is the observation noise covariance.

This is the theoretically optimal weighting between new information and prior belief, and it is the precise thing that Mamba-style selection lacks.

### Deriving the Kalman-Optimal Model

KOSS derives a continuous-time state space model that inherits these Kalman-optimal properties.

KOSS starts with the classic Kalman update formula a Kalman filter uses to update its memory.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/93ef7ff7-67b7-4c43-b406-fca21a4826c2/Continuous-time_SSM_KOSS__1_.png)

By assuming the Kalman gain changes slowly over time (which it does), the researchers derived a clean, continuous-time equation for the SSM:

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/520cffcf-0cde-4a10-96bb-00e7be1fab6f/image.png)

Matrices A and B are defined in terms of A, the Kalman gain K(t), and C, a matrix to map from state to output.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/98fcc3a0-489b-4cdb-8682-49c03d4aab7c/image.png)

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/063d2fc8-1925-40cf-8225-62c9d029c5c9/image.png)

The update equation is made up of 3 terms, determining how the model’s memory changes over time:

1. State Transition: How the past memory carries forward.
2. Input Injection: How new data is added.
3. Derivative Term: A term that adjusts the update based on how quickly the input data is changing.

The key point is that the matrices A(t) and B(t) are fully determined by the Kalman gain K(t) and a base matrix A.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/d96b2551-0143-4d96-8e6b-a854a8458505/Continuous-time_SSM_KOSS.png)

The entire system is therefore governed by a single quantity: the Kalman gain. It means the model’s behavior is constrained to be consistent with optimal estimation.

When translated into code for practical use, the discretized form is:

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/d520b1f8-6e1b-4c87-aaca-640e5174cb4a/image.png)

Here, xt is an estimate of the input derivative at time t.

### Innovation-Driven Selectivity

So, how does KOSS actually calculate that crucial Kalman gain without doing massive, slow calculations at every single step?

It uses a nice shortcut based on the innovation of the Kalman filter.

Instead of solving an expensive matrix equation at every step, KOSS intuitively estimates the Kalman gain by looking at how “surprised” the model is by each new input.

- A large surprise (large innovation) means the model's prior was wrong, so the model updates aggressively with a high gain.
- A small surprise (small innovation) means the prior was good, so the model keeps a low gain and trusts its state.

This is the Innovation-Driven Selectivity (IDS) mechanism.

By considering innovation in selection, the model implicitly has access to context information.

The IDS is the primary reason KOSS is able to outperform models like Mamba with a “closed-loop” advantage. Well, what exactly is “closed-loop”?

Mamba uses an “open-loop” selection mechanism, meaning the weights depend only on the current (new) input with no feedback from the latent state.

KOSS uses “closed-loop” selection, since the innovation feeds the latent state back into the selection weights, creating a context-aware feedback-regulated system analogous to the feedback loop in a physical control system.

With the closed-loop mechanism, distractor tokens are able to be ignored effectively since even if a distractor token superficially resembles a genuine signal, it will produce a different innovation if the latent state has already learned the statistical pattern of genuine tokens.

### The Spectral Differentiation Unit

The KOSS formulation requires estimating how fast the input is changing at each timestep, specifically the input derivative. However, calculating this derivative is quite tricky.

The most obvious approach is finite differences, which just subtracts consecutive values.

But sequences in machine learning are rarely clean. They carry noise at every timescale. Finite differences amplify high-frequency noise dramatically, turning a slightly noisy signal into a chaotic one. So naive differentiation fails here.

The paper proposes the Spectral Differentiation Unit (SDU), which takes a different approach: we estimate derivatives in the frequency domain rather than the time domain.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/13038073-3910-4591-9ef6-589140d1bb62/Spectral_Differentiation_Unit__SDU___KOSS__1_.png)

Here’s how it works:

- They use Fourier transforms (Fast Fourier Transform) to decompose the signal into a mixture of sine waves at different frequencies.
- They differentiate the sine wave, which simply scales it by its frequency.
- They apply a soft spectral mask to dampen high-frequency waves (which are most likely random noise).
- They finally transform the waves back into a regular timeline (Inverse Fast Fourier Transform).

In practice, the steps are modeled by the below equation:

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/f6ebab4e-7a21-4af3-9ab0-db0ea364973b/image.png)

Here, ωk=(2πk)/(NΔt) is the angular frequency of the k-th component. SDU has three concrete advantages over finite differences.

The first advantage is a global receptive field. The SDU uses the entire sequence to estimate the derivative at any point, so it captures slow trends and long-period oscillations.

The second is a “natural” noise suppression. High-frequency components, which are usually noise, can be softly damped before the transform, so the derivative estimate reflects pure signal.

The last advantage is computational efficiency. Through the Fast Fourier Transform, this runs in O(N log N), faster than any attention-based alternative over the same sequence length.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/43a79ab5-4205-4bb1-a2fb-4a925307cf7a/image.png)

KOSS layer from the original paper

### Empirical Performance

The experimental evidence for KOSS is comprehensive. Across nine long-term forecasting benchmarks (traffic, electricity, weather, etc.), KOSS consistently outperforms all baselines, including Mamba variants, Transformers, and MLP-based approaches.

- In terms of accuracy, KOSS beat every baseline model. It improved the mean squared error by roughly 15-30% across benchmarks, with a 36% improvement on traffic prediction.
- KOSS’ accuracy holds steady or actually improves the further out it looks across prediction horizons since the Kalman filter slowly calibrates itself to the noise levels of the data. This is an impressive finding; typically, other models get worse as they get far.
- The model only has 0.2 million parameters, significantly fewer than the 1.17 million parameter transformer baseline, but modestly more than Mamba at 0.07 million. Its clever segment-by-segment architecture uses up to 6x less GPU memory than Mamba during training.

### What This Means for Long-Sequence Modeling

The broader significance of KOSS is in the architectural blueprint.

It demonstrates that estimation theory, developed over decades for linear, Gaussian, low-dimensional systems, can be productively transferred into deep learning for high-dimensional, nonlinear sequence modeling.

The key enablers are:

- The innovation-driven reparameterization that makes the Kalman gain learnable from data;
- The Spectral Differentiation Unit (SDU) that provides stable derivatives;
- And the segment-wise scan that reconciles the resulting state-dependent coupling with hardware-efficient computation.

KOSS suggests that the next generation of sequence models may not come from scaling existing architectures but instead from involving the mathematics of classical estimation and letting it reshape architectures from the ground up.

---

Here’s a special thanks to our [biggest supporters](https://www.aibutsimple.com/upgrade?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=kalman-state-space-models-simply-explained):

- Sushant Waidande
- Sai Krishna Pashapu (AI Developer, BTech, [LinkedIn](https://www.linkedin.com/in/sai-krishna-pashapu-997358283/?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=kalman-state-space-models-simply-explained), [Github](https://github.com/saikrishna-pashapu?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=kalman-state-space-models-simply-explained))
- Bob Boncella
- Demet Yazici

---

If you enjoy our content, consider supporting us so we can keep doing what we do. Please share this with a friend!

Want to reach 50000+ ML engineers? Let’s work together:

[WORK WITH US 📈](https://tally.so/r/A7Dgdk?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=kalman-state-space-models-simply-explained)

Feedback or inquiries? Send us an email at [edwin@aibutsimple.com](mailto:edwin@aibutsimple.com).

---

That’s it for this week’s issue of AI, but simple. See you next week!

—AI, but simple team