What if I told you that the same recursive algorithm that steered Apollo 11 to the moon, that helps neuroscientists decode signals from the visual cortex, and that runs inside every smartphone GPS - is, with a one-line reparameterisation, a linear-attention layer for modern sequence models?

That algorithm is the Kalman filter: sixty-five years old, trusted everywhere engineers need to infer the hidden state of a noisy world from partial observations of it. This three-part series shows how it becomes Kalman Linear Attention (KLA) — a drop-in replacement for softmax attention with linear-time inference, parallel training, and calibrated uncertainty for free.

This first post is the filter itself: the minimum working mental model — no prerequisites beyond basic linear algebra and probability. Already fluent in predict-update? Skip ahead to Part 2, where the same recursion turns into an attention layer.

The journey. §1 builds intuition · §2 sets up the stochastic SSM in which filtering / state estimation is performed · §3 states the problem the KF solves · §4 when and why is it the “best” of all estimators? :) · §5 sketches the derivation of the recursion · §6 watches one work on a James Bond 🕶️ scenario · §7 points to further references.

1. What Is a Kalman Filter?

A Kalman filter is a state observer / estimator - a recursive estimator that infers the latent state (the hidden quantities the world keeps from you) of a (stochastic) black-box physical system from a stream of noisy, partial observations of it. The eye in Figure 1 is more than a metaphor: mathematically, a state observer / estimator is a piece of code that consumes observations $y_{1:t}$ and produces an estimate $\hat x_t$ of “what the world looks like, given what I’ve seen so far”. The Kalman filter is one specific - and very famous - observer: the optimal linear one for a particular class of models.

Figure 1. The unknown system + observer frame the rest of the post lives in. The top box is the world - possibly non-linear, possibly stochastic, with hidden state $x_t$. The bottom box is the observer (the eye is its job: looking at the world through the only window it has, $y_t$). The Kalman filter is the specific observer we'll build.

The KF’s inference equations are derived from a particular linear choice for the system functions $f$ and $g$ above, and a Gaussian white-noise assumption (sometimes weaker - see below). We’ll lay out that model in §2 next.

2. Linear Noisy State-Space Models

We have an unknown system and want to infer its underlying state in real time from a stream of observations. To get anywhere we have to commit to a simplified generative model of how the world produces those observations - with the simplifications chosen so that

1. inference is tractable in real time for the use case (constant compute and memory per step), and
2. the modelling errors and uncertainties introduced by those simplifications are absorbed into explicit noise terms - so the filter always knows what it doesn’t know.

Concretely, the Kalman filter’s simplification is exactly the one previewed at the end of §1: take the transition and observation functions $f, g$ to be linear, and let the noise terms $\omega_t, \gamma_t$ absorb every unmodelled effect. Spelled out, that gives the following linear state-space generative model.

1. Linear dynamics. The latent state evolves linearly in time: \[x_{t+1} \;=\; A_t\, x_t \;+\; B_t\, u_t \tag{2.1}\] where $u_t$ is an optional control input. In unactuated settings - the most relevant for language modelling used in KLAKLA = Kalman Linear Attention, the technique built in Part 2 - a Kalman recursion wrapped in a deep encoder with all four matrices learned. - $B_t u_t$ is absent. 2. Process noise ($\omega_t$) absorbs everything else (mild non-linearities, model errors, external disturbances): \[x_{t+1} \;=\; A_t\, x_t \;+\; B_t\, u_t \;+\; \omega_t \tag{2.2}\] 3. Linear observation model. What you see is a linear function of the state: \[y_t \;=\; C_t\, x_t \tag{2.3}\] 4. Observation noise ($\gamma_t$). Sensors are noisy - same trick: \[y_t \;=\; C_t\, x_t \;+\; \gamma_t \tag{2.4}\] 5. White noise. $\omega_t$ and $\gamma_t$ are zero-mean, uncorrelated across time, with finite (constant) covariance: \[\Sigma^{\mathrm{proc}} = \mathbb{E}\!\left[\omega_t \omega_t^{\top}\right], \quad \Sigma^{\mathrm{obs}} = \mathbb{E}\!\left[\gamma_t \gamma_t^{\top}\right] \tag{2.5}\]

Figure 2. The five assumptions drawn as a graphical model. Latents $x_t$ form a hidden chain (top); each emits a noisy observation $y_t$ (bottom). Shown is the Gaussian special case - the typical textbook form. Under just white noise (per the callout below), the same equations describe only the first two moments of the conditionals.

🟦 On the Myth of Gaussianity Being Necessary. Notice that Gaussianity was nowhere on the assumption list above - that absence is deliberate, not an oversight. The PGM in Figure 2 draws the Gaussian special case only because it makes the textbook Bayesian derivation closed-form; Kalman’s 1960 derivation itself never requires it. Two precise takeaways:

1. Gaussianity isn’t needed to write down or run the Kalman recursion. Kalman’s derivation leaves $p(x_t \mid x_{t-1})$ and $p(y_t \mid x_t)$ as arbitrary, unspecified distributions and characterises them only through their first two conditional moments - the conditional expectations $\mathbb{E}[x_t \mid x_{t-1}] = A_t x_{t-1}$, $\mathbb{E}[y_t \mid x_t] = C_t x_t$ and the conditional covariances $\mathrm{Cov}[x_t \mid x_{t-1}] = \Sigma^{\mathrm{proc}}_t$, $\mathrm{Cov}[y_t \mid x_t] = \Sigma^{\mathrm{obs}}_t$ - with the densities $p(\,\cdot \mid \cdot\,)$ themselves never committed to a parametric form. The closed-form recursive estimator is derived from these alone, via optimisation - no Gaussian density ever invoked.
2. Gaussianity sharpens the optimality claim - it doesn’t unlock the algorithm. Without it, the recursion is the best linear-MSE estimator (no linear estimator beats it; nonlinear ones in principle could). With Gaussianity - or just symmetry+unimodality of the posterior - the same equations deliver the true conditional mean $\mathbb{E}[x_t \mid y_{1:t}]$, which is the best estimator among all reasonable estimators under MSE. Same equations, strictly stronger guarantee. Kalman 1960 lays out the full hierarchy of optimality theorems spanning these (loss, distribution) regimes; we work through them in §4.

2.1. Is the Linear Assumption Too Restrictive?

Looking at that list it’s tempting to say “linear dynamics + linear observations? My problem isn’t like that at all.” In practice the filter is surprisingly robust for two reasons.

1. The noise term is a load-bearing fudge factor. Anything you didn’t model - slow drifts, sensor biases, mild non-linearities - gets absorbed into $\omega_t$ and $\gamma_t$. The covariances $\Sigma^{\mathrm{proc}}$ and $\Sigma^{\mathrm{obs}}$ are tuning knobs that tell the filter how much to trust its dynamics versus its sensors.

2. A wrong-but-good-enough model with reasonable noise estimates can still produce excellent state estimates - online, recursively, in constant memory. The filter doesn’t need to be the truth; it needs to be useful.

If you want richer expressivity - non-linear latent dynamics, learned features, hierarchical structure - replace the handcrafted matrices with neural networks: process incoming sensory signals through a neural feature extractor and do filtering in a non-linear latent space. We’ll explore that path in Part 2: from the linear KF to Kalman Linear Attention (KLA).

3. Toward the KF Update Equations: State Estimation via Optimization (and its Inference Dual)

Earlier we framed the KF as a state observer - a black-box consumer of $y_{1:t}$ that emits a state estimate $\hat x_t$. But how does that estimation actually work? The KF update equations can be seen as solutions to one well-posed problem with two equivalent statements: either (i) solving an optimisation objective with a specific loss $L$, or (ii) doing probabilistic inference on a specific directed graphical model - the Linear Gaussian State Space Model (LGSSM) PGM drawn earlier (Figure 2, the Five Assumptions graphical model). We write both statements here. §5 walks through how the recursion is derived from either one (optional reading), and §5.1 (The Predict–Update Cycle) shows what the recursion does at each step.

Statement A - optimisation. Given $y_{1:t}$, find the estimator that minimises an expected loss

Statement B - inference (the dual). Given the same $y_{1:t}$, compute the posterior over the latent state

and report a summary of it - by default, its mean.

🟦 (A) and (B) Coincide on the LGSSM but importantly not the only case where this can coincide. Kalman 1960 pins down the precise conditions under which the optimisation statement (3.1) and the inference statement (3.2) yield the same answer - illustrated in Figure 3, with the broader duality map detailed in the toggle below (optional).

🟦 The Take-Away for the Rest of This Post. The (A) ↔ (B) duality isn’t a hand-wavy observation - Kalman 1960 establishes it rigorously via a small hierarchy of optimality theorems (Thm 1, 1-a, 2 - laid out in §4) that hold across surprisingly general regimes of loss and distribution. A practical consequence: because the two sides of the duality provably coincide, the same Kalman recursion can be re-derived from the native toolkit of whichever field one is most fluent in - optimisation (orthogonal projection in $L^2$$L^2$ is the Hilbert space of square-integrable random variables - random variables $X$ with $\mathbb{E}[X^2] < \infty$ - equipped with inner product $\langle X, Y \rangle = \mathbb{E}[XY]$ and norm $\|X\| = \sqrt{\mathbb{E}[X^2]}$. The norm-squared is precisely the MSE, so minimising MSE = minimising squared distance in $L^2$, and the best linear estimator is the orthogonal projection of $x_t$ onto the linear span of past observations $\{y_1, \ldots, y_t\} \subset L^2$. Not to be confused with $\ell_2$ (lowercase $\ell$), which in ML usually refers to finite-dimensional weight regularisation / squared loss.), probabilistic inference / graphical models (message passing on the LGSSM PGM similar to HMMs), information geometry (KL projection onto a Gaussian manifold), and so on. In every case the target object is the same: $\mathbb{E}[x_t \mid y_{1:t}]$.

(Optional / Advanced) Mapping the (A. Optimization) ↔ (B.Inference) duality region - where does the Kalman recursion live? →

The precise duality conditions. Kalman 1960 lays this out in two theorems: under MSE, no distributional assumption on the posterior is needed (Thm 1-a); under broader losses, symmetry+unimodality of the posterior suffices (Thm 1). Spelling out the equality:

\[\underbrace{\arg\min_{\hat{x}}\;\mathbb{E}\!\left[\,L(x_t-\hat{x}) \,\big|\, y_{1:t}\,\right]}_{\text{optimum of (A)}}\;\;=\;\;\underbrace{\mathbb{E}\!\left[\,x_t \,\big|\, y_{1:t}\,\right]}_{\text{mean of (B)}} \tag{3.3}\]

holds in two cases (Figure 3 below):

• Squared-loss case (Thm 1-a). For $L(\varepsilon)=\varepsilon^{2}$ (MSE), the equality holds for any distribution - no symmetry or unimodality needed. $\mathbb{E}[x_t \mid y_{1:t}]$ is simply the $L^{2}$-projection of $x_t$ onto the observations.
• General-loss case (Thm 1). For any “reasonable” loss $L$ - symmetric, non-decreasing in $\lvert\varepsilon\rvert$, e.g., $\lvert\varepsilon\rvert$, $\varepsilon^{4}$, $1-e^{-\varepsilon^{2}}$ - the equality holds provided the posterior is symmetric and unimodal. Gaussian is the canonical instance (symmetric and unimodal by construction), and the LGSSM gives us this for free.

_{Figure 3. Each Kalman 1960 theorem picks out a rectangle in the $(\mathcal{L}, \mathcal{D})$ plane on which $\mathbb{E}[x_t \mid y_{1:t}]$ is the optimum of (A).
Thm 1-a = horizontal strip (MSE $\times$ any distribution); Thm 1 = vertical strip (any reasonable loss $\times$ sym+uni).
They cross at MSE $\times$ sym+uni - Thm 1's stronger all-loss claim dominates there.
Gaussian (and hence the LGSSM) lives in the intersection.}

Where the Kalman recursion sits on this map. The recursion we derive in §5 is constrained to the LGSSM corner: linear dynamics $+$ linear observations $+$ MSE loss $+$ Gaussian (or just sym+uni) noise. That corner sits inside the intersection cell of the strips diagram - both theorems apply, and the conditional mean is optimal under any reasonable loss. The wider map matters because it tells us the duality (A) ↔ (B) holds far beyond the LGSSM corner - but practical, recursive, closed-form algorithms only exist inside it. KLA in Part 2 will be a controlled departure from the corner, via data-dependent dynamics.

Why the dual view is useful. Holding (A) and (B) as two faces of the same problem matters chiefly when you want to derive the recursion from scratch (as sketched in §5) - different techniques (loss-minimization, posterior-update, orthogonal projection) start from different sides of the duality and land on the same equations. For just using the recursion, the intuition alone is enough: state estimation can be cast either way, and the two views coincide on the LGSSM.

4. Best Linear Estimator under MSE Loss

Suppose we restrict ourselves to linear estimators - i.e. estimators that are a linear function of past observations $(y_{1:t})$,

where each $\alpha_s$ is a real-valued (matrix-valued in general) coefficient applied to past observation $y_s$ - one matrix per time index. The $\alpha_s$ aren’t fixed in advance; they are picked by the estimator / optimisation algorithm itself, as whatever values minimise the optimisation objective in eq (3.1) - equivalently, whatever values realise the inference summary in eq (3.2). Theorem 2 below then characterises that optimal pick - it turns out to coincide with the orthogonal-projection coefficients (see the toggle).

Kalman’s Theorem 2 says that, within this class of linear estimators, the optimum under squared (MSE) loss is achieved by the Kalman recursion - and, crucially, no distributional assumption on the noise processes $\omega_t, \gamma_t$ is needed for this to hold. Just second moments. No Gaussianity, no symmetry, no unimodality. The Kalman filter is therefore the gold-standard linear-MSE estimator for any linear state-space model: nonlinear estimators (particle filters, nonlinear smoothers) can in principle do better, but no linear one can.

🟦 Relevance to modern linear RNN / SSM literature. Linear RNNs and linear SSMs are - by construction - linear functions of their input history. Thm 2 says we cannot theoretically do better than the Kalman recursion in an MSE sense within that class. The Kalman filter is the gold-standard linear MSE estimator for any linear state-space model. Beating it requires breaking linearity-in-input - via nonlinear smoothers, particle filters, attention, or data-dependent dynamics (selective SSMs like Mamba and KLA - exactly Part 2’s story).

Under Gaussianity: even stronger. Kalman’s own commentary on Thm 2 notes that Gaussianity broadens the result further. If the underlying process is in fact Gaussian (more generally: symmetric+unimodal posteriors - see the case-split inside the toggle in §3), Thm 1 promotes the Kalman estimator from “best linear under MSE” to “best estimator overall, under every reasonable loss.” This isn’t irrelevant for what follows - there are KLA / SSM applications where the data are at least approximately Gaussian, and you get the stronger all-loss guarantee for free. And when Gaussianity doesn’t hold, Thm 2 alone - best linear MSE estimator - is already enough to make the Kalman recursion a very strong choice.

5. Explaination and Derivation of the Update Equations

As a headsup only understanding what the update equations is a necessary primer for KLA in part 2, the full on derivations are not required and are discussed only for completeness.

As discussed in §3, the inference ↔ optimisation duality on the LGSSM means the same Kalman recursion can be derived from several different toolboxes - each starting from a different side of the duality and landing on the same equations. We’ll call these points of view (POVs).

Regardless of which POV one takes, every derivation of the KF threads through two canonical steps:

1. Prediction (a.k.a. time update) - given past observations $y_{1:t-1}$, what is my best guess for the system state at the next time step $t$, before seeing the new observation $y_t$?
2. Update (a.k.a. observation update) - now fold the fresh observation $y_t$ into that prediction, to get the best guess for the system state at time $t$ given all observations $y_{1:t}$ so far.

The toggle below sketches two such POVs - one via optimisation, one via probabilistic inference on the LGSSM PGM - showing how each threads through these two steps and lands on the same recursion. The full derivations are beyond the scope of this primer and may be covered in separate posts; open the toggle only if you’d like intuition for the different tools one can use to derive the KF (and adapt them to your own SSM variants beyond the simple linear-Markovian case - switching models, hierarchical models, etc.). The Predict–Update cycle in §5.1 below is what you actually need to follow Part 2.

What you actually need to carry into Part 2. You will not rederive any of this. What matters is: at every step the KF emits a pair $(\mu_t, P_t)$ - posterior mean and covariance - and these obey a linear recursion in those quantities, equivalently in the natural parameters $(\Lambda_t, \eta_t)$. Some linear algebra turns the predict–update cycle into the closed-form Kalman equations; which derivation got you there (orthogonal projection, Bayes’ rule, exponential-family conjugacy) is a matter of taste. Part 2 holds the recursion structure fixed and reparameterises it for parallel computation - it does not rederive it.

5.1. The Predict-Update Cycle

To make what the filter does at each step concrete, here is what happens to its belief between two observations:

Figure 5. The predict–update cycle, two views.

Left: the predict step can be cast as either solving an optimisation problem (argmin of expected squared error from past observations) or doing inference. For the NLP / HMM reader, that inference is just the marginalisation substep of the forward algorithm - known in continuous state-spaces as the Chapman–Kolmogorov equation. The update step is similarly either an optimisation (argmin after folding in the latest $y_t$) or inference - the conditioning substep of the forward algorithm, i.e. Bayes’ rule applied to $y_t$. On the LGSSM both views collapse to the same closed-form recursion shown below each.

Right: visual intuition on how the belief evolves on a 1D state axis. The dot is the mean (first moment); the shaded region around it is the variance (second moment) — only these two moments are committed, no Gaussianity required (matching Thm 2 of §4). posterior at $\textcolor{#2E7D32}{t-1}$ → predict widens it to the prior at $\textcolor{#C66A00}{t}$ (variance grows by additive process noise $Q$) → update folds in the observation $\textcolor{#D32F2F}{y_t}$, giving a tighter posterior at $\textcolor{#2E7D32}{t}$. Predict always widens uncertainty; update always tightens it — or, in the worst case (missing observation $\equiv$ infinite observation variance / zero precision), leaves it equal. Update never makes things worse.*

For reference - the closed-form LGSSM recursion shown inside Figure 5, written out for copy-paste convenience.

Predict step - propagate the previous posterior through the dynamics:

Update step - fold in the new observation $y_t$ via the Kalman gain $K_t$:

6. A Classical Kalman Filter in Action

Enough abstraction. Let’s see one work.

🎬 The setup. James Bond is on vacation, having a private call in a closed 1-D corridor. His security agency wants to track his movements for safety without compromising privacy - they only have an inexpensive, noisy position sensor on his phone. The agency uses a Kalman filter to fuse the sequence of sparse, noisy sensor readings and estimate / predict in real time Bond’s position and velocity.

Figure 6. The full pipeline in motion. Watch the orange ghost overshoot during predict steps, the green posterior snap back at every ping, and the velocity panel quietly track the truth even though it never receives a single direct observation.

6.1. What to Watch in the Animation

A quick tour of the four panels, top to bottom:

Top - graphical model + predict/update. Kalman filtering viewed as posterior inference in an SSM. The predict step rolls the latent $x$ forward through the causal dynamics $(A, \Sigma^{\mathrm{proc}})$ to give the prior $p(x_t \mid y_{1:t-1})$; the update step performs a Bayesian inversion - the green dashed back-arrow $y_t \to x_t$ in the PGM - to fold in the new noisy observation and give the posterior $p(x_t \mid y_{1:t})$. When no ping arrives at a step, we simply carry the prior forward.

Second row - position over time. True position $p_t$ in blue; noisy position sensor data ($y_t$) as red dots; filter mean $\hat{\mu}_p$ in orange, with the $\pm 2\sigma$ uncertainty band shaded around it.

Third row - velocity over time. Same colour scheme for the latent velocity, except there are no red dots (🔴) - no velocity sensor exists. The filter must infer velocity purely from the pattern of incoming position sensor data.

Bottom - 1-D corridor (spatial view). Bond (solid blue 🧍) paces the corridor in real time; the dashed ghost is the filter’s belief - orange (🧍) during predict (the prior), green (🧍) during update (the posterior). The footprint underfoot is the same $\pm 2\sigma$ band as in the panels above. Each ping (🛜) appears as red wifi arcs above the corridor and a red triangle on the floor at the noisy reading.

Two things to notice - and to expect in most Kalman applications:

1. Between pings (🛜), the filter coasts on its predict step (🧍) and tends to drift from the truth. With no new evidence, the prior can only roll forward through the assumed dynamics; the estimate becomes inaccurate and the uncertainty band grows under the process noise $\Sigma^{\mathrm{proc}}$. The next ping then grounds the estimate, snapping it back toward truth and shrinking the band.

2. Hidden quantities can be recovered reliably, even when direct sensing is impossible. Velocity here is a tame example - in practice the same machinery recovers things like neuronal activity in the brain from indirect biosignals, or atmospheric state from sparse weather stations. The Kalman filter pieces them together from the pattern of related observations, in real time and at constant compute per step.

6.2. Designing the Generative Model

Now let’s pop the hood. In a classical application of Kalman filtering, the modelling effort is front-loaded: you spend most of your time designing the generative model - the transition matrix $A$, the observation matrix $H$, and the noise covariances $\Sigma^{\mathrm{proc}}$ and $\Sigma^{\mathrm{obs}}$ - and only then turn the predict-update recursion of §5 loose on the incoming data. Here are the four design choices that produce the Bond scenario, with the trade-offs and intuition behind each.

Step 1 - pick the latent state. We track

position and velocity, even though no sensor ever reads velocity. Why include it? Position alone tells you where Bond is, not where he is going - without velocity, the filter has no way to extrapolate between pings. The latent state is exactly the set of quantities you’d need to roll the world forward by one step.

Step 2 - pick the transition matrix $A$. The agency picks the simplest reasonable dynamics - constant velocity:

In words: “position drifts forward by velocity $\times \Delta t$; velocity stays the same.” This model is deliberately wrong - Bond’s true motion is a sinusoid, not a straight line. That’s the universal practical lesson: in any real application your dynamics is always a simplification, and the role of $\Sigma^{\mathrm{proc}}$ is to tell the filter how wrong.

Step 3 - pick the observation matrix $C$. The cheap phone sensor only reads position, so

$C$ literally selects “the position component of $x_t$”. Richer sensors (e.g. position and heading) would just give $C$ more rows.

Step 4 - pick the noise covariances. $\Sigma^{\mathrm{proc}}$ and $\Sigma^{\mathrm{obs}}$ are the filter’s two tuning knobs - they encode trust rather than measure anything physical:

The small $(1,1)$ entry of $\Sigma^{\mathrm{proc}}$ says “position barely drifts on its own” - it’s mostly determined by velocity. The larger $(2,2)$ entry says “I expect velocity to wander by $\approx \pm 0.22$ per step” - wide enough to absorb the unmodelled acceleration of the sinusoidal truth. The $(0.20)^2$ on the sensor side says “the phone sensor is accurate to about $\pm 20$ cm.” What ultimately matters for the Kalman gain $K_t$ is the ratio $\Sigma^{\mathrm{proc}} / \Sigma^{\mathrm{obs}}$ - not the absolute scales.

With $A, H, \Sigma^{\mathrm{proc}}, \Sigma^{\mathrm{obs}}$ in hand, the rest is mechanical: feed the incoming pings $y_2, y_4, \dots$ from the noisy position sensor into the predict-update recursion of §5, and read off the posterior mean $\mu_t$ and covariance $P_t$ at every step. The remarkable thing is that even though velocity is never directly observed, the recursion produces a perfectly reasonable real-time estimate of it - and at constant compute per step, because the running posterior $(\mu_t, P_t)$ is a sufficient summary of everything the filter has seen so far.

💡 What changes in KLA (Part 2). KLA reuses the recursion of §5 but departs from the classical setup above in two structural ways.

1. Non-linear feature lift. The “noisy observations” $y$ - for KLA, input tokens - are first pushed through a deep encoder before they enter the recursion. The Kalman update still runs in a linear-Gaussian latent space, but that latent now lives in the encoder’s feature space rather than the raw observation space.
2. All four quantities are learned, not hand-designed. $A, H, \Sigma^{\mathrm{proc}}, \Sigma^{\mathrm{obs}}$ become differentiable parameters fit by gradient descent on a downstream task - no constant-velocity assumption, no hand-tuned noise scales, no manually-picked $C = (1\;\;0)$.

The recursion itself is unchanged; what changes is who picks the matrices.

If you’d like to play with the simulation - change the noise covariances, the obs rate, the truth itself - see the companion Jupyter notebook (kf_tutorial.ipynb in the GitHub repo), which builds the filter step by step with running plots.

7. Going Further

The actual derivation of the predict–update equations is beyond the scope of this primer; both views (Bayesian and orthogonal-projection) are spelled out in detail in the references below.

• Shaj, V. et al. (2026). Kalman Linear Attention: Parallel Bayesian Filtering for Efficient Language Modelling and State Tracking. ICML 2026. The KLA paper this primer series is built around - reframes linear attention as a learnable Kalman recursion with closed-form predict-update steps and linear-time parallel training.
• Kalman, R. E. (1960). A New Approach to Linear Filtering and Prediction Problems. The original paper - surprisingly readable, and worth reading for the geometric (Hilbert-space) derivation that modern ML treatments often skip.
• Bishop, G. & Welch, G. (2001). An Introduction to the Kalman Filter. The classic SIGGRAPH course note - a compact, practitioner-oriented walkthrough that complements Kalman’s original paper.
• Bzarg. How a Kalman filter works, in pictures. The clearest visual intuition I know of for the update step (overlapping Gaussians).
• Eric Xing. Sequential Models (lecture). A graduate-level lecture connecting HMMs, Kalman filters, and the broader family of sequential latent-variable models - useful for placing the KF inside the wider PGM landscape.

In Part 2, we’ll see how the predict–update recursion can be reframed as a linear-attention layer, and how that leads naturally to Kalman Linear Attention (KLA) - our drop-in probabilistic primitive for sequence modelling that retains all the parallel-training goodness of modern transformers but adds principled uncertainty estimates and richer state-tracking expressivity.

Probabilistic Graphical Models (PGMs) are powerful tools for representing relationships between random variables using graphs (directed or undirected). Once we come up with a representation of the stochastic phenomenon we wish to model, PGMs provide a consistent and flexible framework to devise principled objectives, set up models that reflect the causal structure in the world, and allow a common set of inference methods to be deployed against a broad range of problem domains.

This blog will try to summarize how Reinforcement Learning can be brought under the PGM framework which allows us to transfer policy search from an optimization point of view to an inference point of view. Also, we will see how this framework allows us to recover “soft” versions of Bellman backup equations from classical RL.

Premise :

Human or animal behaviour is not often perfectly optimal but approximately optimal. For example, an animal/human being whose only goal is to move from the start position to the goal position as shown in figure 1 can choose any choose either perfectly optimal/hard optimal (green) trajectory or suboptimal (blue) trajectory but typically would avoid the red trajectories which specify bad behaviour.

We need a probabilistic framework for this stochastic phenomenon of goal-directed behaviour. The framework, in addition to giving a higher probability to perfectly optimal behaviour, should also give a non-zero probability to suboptimal behaviour. Similarly, a near-zero probability should be assigned to bad behaviour that misses reaching the goal.

In machine learning whenever we have a stochastic phenomenon, we usually come up with a probabilistic graphical model (PGM) based on the observed stochastic phenomenon, such that the samples from this model will look like the observed stochastic phenomenon.

Thus we need a PGM that models optimal decision-making. Let’s first model a PGM that models the relationship between states, actions and next states as in figure 2.

This PGM in figure 2, commonly known as a state space model (SSM), can represent physically consistent trajectories but for control/decision making we need a notion of cost/reward/optimality. Thus we modify the PGM to introduce additional optimality variables into the PGM, which is defined as $p(O_t=1 \mid s_t,a_t) = exp(r(s_t,a_t))$.

Thus $O_t$ is defined as a binary random variable which indicates if the behaviour at time t was optimal or not. The reason why we define this as the exponential of the reward function is to make a convenient distinction between optimal (best way to reach a goal), sub-optimal (reaches the goal but not optimally) and bad behaviour (fails to reach the goal). This will become more apparent in the next section.

Motivation For Why Control As Inference Is An Interesting Paradigm

1. Convenient Way For Representing and Sampling Optimal and Suboptimal Behaviour

Let’s derive a mathematical expression for the probability of a trajectory (state action sequence) given the observed optimality variables. We will see that based on the PGM we can derive an expression where optimal/sub-optimal trajectories can be conveniently represented.

As we can see, those trajectories which are physically consistent and optimal (in terms of long-term rewards) have a higher probability mass. Additionally, a suboptimal trajectory with a slightly lesser reward can also be modelled/sampled using this graphical model framework which is important in several settings including inverse reinforcement learning.

Note: One could ignore the action prior or assume a uniform action prior to further simplify the equations.

2. Can use established inference schemes to answer several queries including:

• Policy Search: Given a reward, infer the optimal policy by calculating $p(a_t \mid s_t, O_{t:T})$. Instead of solving the optimization problem, we now can solve the inference problem. This will be discussed in detail in part 2 of this blog. Further, an approximate inference scheme based on variational/optimization-based formulation is discussed in part 3 of this blog.

• Inverse Reinforcement Learning: Given a collection of optimal trajectories, infer the reward and action priors, which is basically an inverse RL question.

  \[\begin{aligned} p\left(\tau, \mathcal{O}_{1: T}, \theta, \phi\right) & \propto\left[p\left(s_{1}\right) \prod_{t=1}^{T} p\left(s_{t+1} \mid s_{t}, a_{t}\right)\right] \times \\ &\quad\quad\quad\quad\quad\quad\exp \left(\sum_{t=1}^{T} r_{\phi}\left(s_{t}, a_{t}\right)+\log p_{\theta}\left(a_{t} \mid s_{t}\right)\right) \\&=\left[p\left(s_{1}\right) \prod_{t=1}^{T} p\left(s_{t+1} \mid s_{t}, a_{t}\right)\right]\times \\ &\quad\quad\quad\quad\quad\quad\operatorname{exp}\left(\sum_{t=1}^{T} \phi^{T} f_{r}\left(s_{t}, a_{t}\right)+\log \theta^{T} f_{p}\left(a_{t} \mid s_{t}\right)\right) \end{aligned}\]

3. Allows to model stochastic behaviour which has several advantages

• Transfer Learning: If we can model multiple ways to solve a particular task, this turns to be relevant for transfer learning in a new setting where the task has to be solved in a slightly different manner.

• Better Exploration Strategies: We will see that the maximum extropy objective that we derive in the 3rd part of this blog series on Policy Search as Variational Inference will provide a natural exploration strategy based on entropy maximization.

The blog is based on the following reference, Levine, 2018.

We will see that the subroutines in the policy search procedure in this graphical model result in “soft” variations of bellman update equations, where the hard max in operation is replaced by a softmax.

Log Sum Exp Trick

A useful “trick” to remember before we jump into control as inference procedure is “Log Sum Exp” trick. The LogSumExp (LSE) (also called RealSoftMax or multivariable softplus) function is defined as the logarithm of the sum of the exponentials of the arguments:

The LSE function is a smooth maximum – a smooth approximation to the maximum function, mainly used by machine learning algorithms. In the following inference procedure, we will replace LSE with max/softmax to derive a soft version of classical RL.

Exact Inference By Recursive Computation Of Backward Messages

At any time t policy search involves computing the posterior over the action $a_t$, given state $s_t$ and the optimality variables $O_{t:T}$, i.e.

Thus policy search involves computing two backward messages \(\color{orange} \beta_{t}\left(\mathbf{s}_{t}\right)\) and \(\color{purple} \beta_{t}\left(\mathbf{s}_{t},\mathbf{a}_{t}\right)\). This is computed via backward messages similar to HMM or Kalman Smoothers as follows:

At the last time step T:

Note on Action Prior: Here $p(a_t \mid s_t)$ is the action prior. Note that it is not conditioned on $O_{1:T}$ in any way, i.e. it does not denote the probability of an optimal action, but simply the prior probability of actions. The PGM for RL as inference doesn’t actually contain this factor, and we can assume that $ p\left(a_t \mid s_t \right)= \frac{1}{ \mid \mathcal{A} \mid} $ for simplicity. That is, it is a constant corresponding to a uniform distribution over the set of actions.

Thought Exercise:

Levine 2018 assumes a uniform action prior and argues that this assumption does not introduce any loss of generality. One could show that any non-uniform action prior $p(a_t \mid s_t)$ can be incorporated into $p(O_t \mid s_t,a_t):=exp(r_1(s_t,a_t))$ via a modified reward function $r_1(s_t,a_t)$.

At any time step t

Intuition / Relationship To Classical RL

In order to get an intuitive meaning of these messages, we do some algebraic manipulation in the log space to get a form similar to bellman backup. The messages in log space are as follows:

One can show that if we replace the $\log \beta(s_t,a_t)$ with Q function and $\log \beta(s_t)$ with the value function we can recover an intuitive relationship resembling bellman backup operator in the deterministic case and an optimistic bellman backup in the stochastic case. Let,

Using Equations 1 and 2, now the messages in the log space look as follows:

Deterministic Case (Regular Bellman Backup)

In the deterministic case since we only have one possibility for transition dynamics, we obtain a backup equation similar to the regular Bellman backup equation.

Stochastic Case (Optimistic Backup)

In the stochastic case we obtain a backup equation that is optimistic as shown below:

The optimistic update occurs because it is largely determined by the max of the next state value, which creates risk-seeking behaviour.

This issue will be mitigated by variational inference discussed in the next part of this blog series.

Which objective does the inference procedure in Exact Inference in Part 2 solve?

The inference procedure discussed in part 2 of this blog series solve the following objective: $\color{red} \text{minimize} \quad D_{\mathrm{KL}}(\color{green} q_\phi(\tau)\color{red} | \color{orange} p(\tau)\color{red}) = \color{red} \text{minimize} \quad D_{\mathrm{KL}}(\color{green} q_\phi(s_{1:T},a_{1:T}) \color{red} | \color{orange} p(s_{1:T},a_{1:T},O_{1:T})\color{red})$

Here the join distribution of optimal trajecories is given as follows:

Which variational distribution to choose ??

Looking at the graphical model for the variational distribution, the joint distribution for \(q(\tau)\) should be \(q(\tau)=q\left(\mathbf{s}_{1}\right) \prod_{t=1}^{T} q\left(\mathbf{s}_{t+1} \mid \mathbf{s}_{t}, \mathbf{a}_{t}\right) \pi\left(\mathbf{a}_{t} \mid \mathbf{s}_{t}\right)\).

Here unlike in the exact inference case, we make an explicit assumption on what part of the graphical model is controllable by agent and what is not. It is reasonable to assume that the transition dynamics is not controllable by the agent and hence we fix \(q\left(\mathbf{s}_{1}\right)=p\left(\mathbf{s}_{1}\right) \text { and } q\left(\mathbf{s}_{t+1} \mid \mathbf{s}_{t}, \mathbf{a}_{t}\right)=p\left(\mathbf{s}_{t+1} \mid \mathbf{s}_{t}, \mathbf{a}_{t}\right)\).

Derivation Of Max Entropy RL Objective

In can be shown that minimizing this optimization objective results in max entorpy reinforcement learning objective as derived below:

Deriving Soft Bellman Equations

We now look at message passing (backward messages) from an optimization point of view. To calculate the backward messages we start from the last time step.

At the last time step T

However, note that here we consider a general scenario where the reward can take any real value, $-\infty < r(s,a) < \infty$ , as opposed to the earlier restriction to be negative or zero. Thus we need to normalize $\exp(\log(r(s_T,a_T)))$, using the normalizing constant $V(s_T)=\int_{\mathbb{A}}\exp(r(s_T,a_T)) da_T$.

Thus we do a little bit more algebraic manipulation to include this normalization constant as follows:

The optimal policy that minimizes this objective is given as :

At any time step t,

The optimal policy that minimizes this objective at any time step t is given as :

This means that, if we fix the dynamics and initial state distribution, and only allow the policy to change, we recover a Bellman backup operator that uses the expected value of the next state, rather than the optimistic estimate we saw in part 2 of the blog series. Thus we avoid the risk seeking behaviour / optimistic bellman backups via the control as inference framework.

We will discuss how this framework is used practically in modern Deep RL alogorithms in the next part of this blog series.

Meta Learning (“Learning To Learn”)

Many problems of interest require rapid inference from small quantities of data. We should adopt learning strategies such that the single/few observations should result in abrupt shifts in behavior(single/few shot learning). When new data is encountered, the models must inefficiently relearn their parameters to adequately incorporate the new information without catastrophic interference.

For example a robot designed to clean nuclear wastes, encounters a wide variety of tasks with many objects. It should be able to amortize their experience from previous learnt skills and improve data efficiency in acquiring new skills rather than Learning each skill from scratch.

Another example would be malware vs clean classification task, where a classification engine has to adapt to the new malware variants being released each day. You would typically have few example per class for a new day(here the new task is classifying between malware and clean files encountered that particular day) and requires meta-learning for “few shot classification”. (This would be an interesting application of meta learning to try out)

Inspired from nature ? This kind of flexible adaptation is a celebrated aspect of human learning (Jankowski et al., 2011), manifesting in settings ranging from motor control (Braun et al., 2009) to the acquisition of abstract concepts (Lake et al., 2015). Generating novel behavior based on inference from a few scraps of information – e.g., inferring the full range of applicability for a new word, heard in only one or two contexts – is something that has remained stubbornly beyond the reach of contemporary machine intelligence.

Why not Deep Learning ? In situations when only a few training examples are presented one-by-one, a straightforward gradient-based solution like Deep Neural Network is to completely re-learn the parameters from the data available at the moment. Such a strategy is prone to poor learning, and/or catastrophic interference. In view of these hazards, non-parametric methods are often considered to be better suited.

General Meta Learning Framework

Meta-learning generally refers to a scenario in which an agent learns at two levels, each associated with different time scales. Rapid learning/Task Specific Learning occurs within a task, for example, when learning to accurately classify within a particular dataset. This learning is guided by knowledge accrued more gradually across tasks, which captures the way in which task structure varies across target domains. Given its two-tiered organization, this form of meta learning is often described as “learning to learn”. This task aganostic learning helps in quick adaptations to new tasks.

This means during the meta-learning we need to provide the learning algorithm an information(encoding) about the task/context in some form, along with the training samples asossiated with each task. So how do we feed data to a meta-learner ?

During meta-learning we provide a support/context $D_{support}$ along with the test sample.

What is $D_{support}$?: $D_{support}$ can be defined as a set of $(x_i,y_i)$ tuples where $x_i$ is input/observation and $y_i$ is the output/reward. $D_{support}$ provides a context(information about task to be performed) with respect to the test sample. $D_{support}$ usually tends to have less number of samples(few shot) per class.

Model Aganostic Meta Learning

The two key features of MAML are,

1) it makes gradient based solutions good at “few shot learning”. 2) it is model aganostic in the sense that it can be applied to any learning algorithm/model(classification, regression and reinforcment learning) that uses a gradient descent based optimization.

Note: There are recent papers that extends maml for UnSupervised and Semi-Supervised tasks.

As discussed earlier this meta learning procedure has 2 stages.

In the task specific update stage, corresponding to each task the model parameters are updated using the task specific dataset.

Update 1 \[ \theta_i^{*}=\theta + \alpha\times \nabla_\theta L_i(\theta,y_i) \]

Next is the meta update stage, where we try to find a genaralist structure among these tasks such that updated parameter is always closer to each of the task specific optimal parameter, $\theta_i$(See Figure 1). In effect the parameter gets updated in the meta update in such a manner that it is one(few) steps away from doing well at each one of the tasks. This is achieved by optimizing for the performance of $f_{\theta_i^*}$ with respect to $\theta$ across the tasks as follows:

Update 2 \[ \begin{aligned} \theta = & \theta + \beta\times\sum_i^NL_i(f_{\color{red}{\theta_i^*}},y_i)
= &\theta + \beta \nabla_\theta \sum_i^NL_i(f_{\color{red}{\theta + \alpha\times \nabla_\theta L_i(\theta,y_i)}},y_i) \end{aligned} \]

Another way of looking at MAML: The meta optimization is performed over the model parameters $\theta$ but the objective is computed using the updated models $f_{\theta_i^*}$. Thus we update $\theta$ such that it learns to get a low classification error/high reward in the next step rather than this step. This makes the model parameter $\theta$ to be more sensitive to changes in tasks such that a small changes in the parameter can produce large improvements in the direction of the task with respect to which it is tuned for.

Does MAML falls in the general meta learning framework ?: Yes. It does. The task specific learning happens in update 1 and meta-learning happens in update 2.

How to train? “You get good at what you practice”. Thus the training time protocol used is same as the testing time protocol, as in O. Vinyal’s(2016). Split each data/demonstrations for each individual task into training and validation pair. Use one of these for Update 1(task specific learning) and the other for Update 2(meta-learning).

An Illustration

Unsupervised Extension

Refer these One-Shot Visual Imitation Learning via Meta-Learning , Semi-Supervised Few-Shot Learning with MAML. I will write on these shortly.

References

1. C. Finn et al. 2017 Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks
2. C. Finn et al. 2017 One-Shot Visual Imitation Learning via Meta-Learning
3. Semi-Supervised Few-Shot Learning with MAML
4. O. Vinyals, 2016, Matching Networks for One Shot Learning

In Kernel theory we assume that learning happens in the RKHS space(Nice space of functions for non-parametric statistics and machine learning) and the theorem that forms the backbone for learning in RKHS is the Representer Theorem.

Before scaring you guys with RKHS and theorems right from the beginning, let me explain two main properties of Kernel methods.

Property 1: Kernel Methods can be thought of as instance-based learners: rather than learning some fixed set of parameters corresponding to the features of their inputs, they instead “remember” the \(i\)-th training example \(\mathbf{(x_i,y_i)}\) and learn for it a corresponding weight \(w_{i}\). This basically means functions(hyperplanes in SVM /basis functions in KPCA) learnt through Kernel Methods can be represented as a weighted linear combination of the training points and what the algorithm actually “learn” are these weights correspondning to each point. Representer Theorem provides explanation for this. Thus as per theorem the ** regularized risk functional**(basically the objective function of the optimization problem being solved) of any algorithm which is a member of the Kernel methods takes the following general form:

\(\hspace{1cm} \Omega-\text{any monotonically increasing function}\) \(\hspace{1cm} C - \text{cost function}\) \(\hspace{1cm} f - \text{function that we intend to learn}\)

Property 2: Now comes the “Kernel Trick”. Kernel Methods through Kernel Functions allow you to perform the learning as if it were projected to a higher dimensional space, by operating on its original space. You can kernalize an algorithm by reformulating it in such a manner that the computations dependent only on the inner products of the data points than actual data points which then can be replaced by a Kernel Function(following the property of RKHS).

• “Why is it nice?”: Say you have a set of images of 32x32=1024 pixels. However this is not linearly separable. The option you have is to project it on to a higher dimensional space, lets say quadratic, $\mathcal{R}^{1024} \rightarrow \mathcal{R}^{1024x1024}$. Data manupulation in this space is highly expensive. Kernel functions allow us to do computations in the lower dimensional space, but effectively learning in the higher dimensional space.

Few Algorithms

1. Support Vector Machines

2. Kernel Ridge Regression

3. Kernel PCA

Appendix

Reproducing Kernel Hilbert Space(RKHS): is a subspace of the Hilbert space with respect to kernel \(k: \mathcal{X}\times\mathcal{X} \rightarrow \mathcal{R}\) constructed in the following manner.

• Let \(\mathcal{H'} = \\{ k(.x): x\epsilon \mathcal{X} \\}\) be set of kernel functions.

• We construct a vector space \(\mathcal{H}\) using the linear combinations of all kernels functions in set \(\mathcal{H'}\), and any function \(f\epsilon \mathcal{H}\) can be represented as a linear combination of a subset of these kernel functions as \(f = \sum_{i=1}^n \alpha_ik(,x_i)\) for some \(n,x_i \epsilon \mathcal{X}\) and \(\alpha \epsilon \mathcal{R}\).

• the inner product of \(f(.),g(.) \epsilon \mathcal{H}\) has the following definition: \(\langle f,g \rangle = \sum_{i=1}^n \sum_{j=1}^{n'} \alpha_i \beta_j k(x_i,x_j)\). (it can proven that this inner product is well defined)

• this parculiar definition gives rise to the reproducing property of hilbert space. \(\langle f,k(.x) \rangle = \sum_{i=1}^n \alpha_i k(x,x_i) = f(x)\). This shows that kernel is a representer of evaluation(evaluation functional), analogous to Dirac delta functions.

• What is it’s significance? : Enables Kernel Trick. While learning in RKHS, inner-products in our computations can be replaced by kernels \(\langle \phi(x),\phi(y) \rangle = \langle k(.,x),k(.,y) \rangle=k(x,y)\), where \(\phi\) maps \(x\epsilon \mathcal{X}\) to an infinite dimensional space, \(\phi : x \rightarrow k(.,x)\epsilon\mathcal{H}\). This is particularly useful in cases where data is not linearly separable in $ \mathcal{X} $, where tranformations to higher dimensional spaces are necessary and kernel trick avoids us in making this explicit transformations.

Representer Theorem: states that a minimizer \(f^{*}\) of a regularized empirical risk function defined over a Reproducing Kernel Hilbert Space can be represented as a finite linear combination of kernel products evaluated on the input points in the training set data.

• Precise Definition: Let \(\Omega : [0,\infty) \rightarrow \mathcal{R}\) be a strictly a monotonically increasing function, by \(\mathcal{X}\) a set, and \(C : (\mathcal{X} × R^2)^N\) be an arbitrary loss function. Then any $ \mathcal{f} \epsilon $ RKHS \(\mathcal{F}\) minimizing the regularized risk functional

\(\frac{1}{N}\sum_{i=1}^{N}C(y_i,f(x_i)) + \frac{\lambda}{2}\Omega(f)\) admits a representation of the form \(\mathcal{f}(.) = \sum_{i=1}^{N}\alpha_i k_{x_i}\).

• What is it’s significance?: Firstly it gave a genreric definition of optimization objective, that come under the umbrella of kernel methods. Secondly, Representer theorems[Smola et. all 2011] are useful from a practical standpoint because they dramatically simplify the regularized empirical risk minimization problem. In most interesting applications, the search domain \(H_{k}\) for the minimization will be an infinite-dimensional subspace of \(L^{2}({\mathcal {X}})\), and therefore the search (as written) does not admit implementation on finite-memory and finite-precision computers. In contrast, the representation of \(f^{*}(\cdot )\) afforded by a representer theorem reduces the original (infinite-dimensional) minimization problem to a search for the optimal {\displaystyle n} n-dimensional vector of coefficients \(\alpha =( \alpha _{1},...,\alpha _{n})\in \mathbb {R} ^{n}\); \(\alpha\) can then be obtained by applying any standard function minimization algorithm. Consequently, representer theorems provide the theoretical basis for the reduction of the general machine learning problem to algorithms that can actually be implemented on computers in practice.

Tip

References

1. Justin Domke - UMASS Notes
2. Sumitra.S.Nair - IIST Notes
3. NPTEL - P S Shastry -IISc
4. Eric Xing - CMU Slides
5. Wikipedia

Disclaimer

Please feel free to contact me vaisakhs.shaj@gmail.com in case you find any errors/suggestions, I will correct those promptly. You may as well comment below.