Kalman Delta Networks: Uncertainty-aware Associative Memory

Linear Attention Rewritten with a Kalman Filter: Kalman Delta Networks

TL;DR

The fixed-size recurrent memory update of linear attention is reinterpreted as a linear-Gaussian state-space model, and the covariance tracking of the optimal estimator, the Kalman filter, is introduced to decide “how strongly to overwrite” based on uncertainty, yielding Kalman Delta Networks (KDN). DeltaNet, Gated DeltaNet, and KDA reduce to fixed-gain special cases that omit the covariance, and KDN outperforms every SOTA linear recurrent mixer evaluated at 750M/1.3B pretraining in both PPL and average accuracy (source: Abstract, §1).

Key Idea

self-attention keeps every past key-value pair as-is, so the cache grows with context length and the computation becomes quadratic. Linear attention compresses this history into a fixed-size state $S_t \in \mathbb{R}^{d_k \times d_v}$, gaining constant-memory decoding and scan parallelization, but instead inherits the “online memory management problem” (source: §1). Each token must edit the past summary without knowing which associations future queries will require; updates that are too weak leave stale information behind, and ones that are too strong destroy useful memories.

Existing delta-rule models (DeltaNet, Gated DeltaNet, KDA) predict the residual write strength $\beta_t$ directly from the current token embedding. The problem is that they do not track confidence in the memory estimate. A residual on an association confirmed by repeated observations and a residual on a never-seen key should be treated differently even when the residual magnitudes are the same, but existing models have no mechanism to distinguish them (source: §1).

KDN’s core claim is this: view the recurrent associative memory as a linear-Gaussian state-space model of a latent non-stationary key-value map, and propagate not only the memory mean but also its covariance (uncertainty) through a Kalman filter; the write strength is then automatically adjusted by the “accumulated evidence” and the “observation confidence” (source: §3). Uncertainty is precisely the state variable missing from delta-rule models.

Background: The Problem They Address

The progress of the linear attention family is a history of how sophisticatedly “forgetting” and “writing” are modeled (source: §2).

• Linear attention: $S_t = S_{t-1} + k_t v_t^{\top}$ — only additive writing is possible, so old associations cannot be explicitly erased.
• DeltaNet: $S_t = S_{t-1} + \beta_t k_t (v_t - S_{t-1}^{\top} k_t)^{\top}$ — an error-correction rule that reads the memory’s prediction for the current key and writes only the residual. For a repeated unit-norm key with $\beta_t = 1$, it overwrites exactly.
• Gated DeltaNet: adds a scalar decay $\alpha_t$ to exponentially fade neglected associations. But since it is a single scalar, all key channels share the same retention rate.
• KDA: replaces the scalar decay with a diagonal transition $D_t = \text{diag}(\alpha_t)$, allowing a different retention rate per channel. It is the most expressive transition in this family, but the write strength $\beta_t$ is still a scalar gate predicted from the token.

The Mamba family uses the selective SSM recursion $h_t = A_t h_{t-1} + B_t x_t$, directly adding the control input $B_t x_t$. Delta-rule mixers, by contrast, correct the predicted key-value map with a key-conditioned residual (source: §2). These two have been understood through different lenses (state-space dynamics vs. online optimization), and they share a common limitation: “predict the write gain from the token but do not track evidence.”

Research gap: previous models combine transition (forgetting) and write (writing), yet the signal that determines write strength was unrelated to “how strongly this association is backed by repeated observations.” A mechanism that tracks uncertainty as an explicit state was absent (source: §1, §6).

New Approach: Kalman Delta Networks

1) Kalman Associative Memory — the delta rule as a Kalman filter

The authors view the memory $S_t$ as an estimate of a latent associative map $\widetilde{S}_t$. Each token provides a noisy observation $v_t$ along one direction (key $k_t$) of this map (source: §3):

Here $D_t$ is the transition (the process model, i.e., the formal target of forgetting), $W_t \sim \mathcal{N}(0, \Omega_t)$ is the process noise (drift not explained by decay), and $e_t \sim \mathcal{N}(0, r_t I)$ is the observation noise (the part of the token value that is “not a clean memory target”).

Under the linear-Gaussian assumption, the optimal recursive estimator in the minimum mean-squared-error (MMSE) sense is the Kalman filter, whose update takes the following predict-correct structure (source: Proposition 1, §3):

where $\widehat{S}_t = D_t S_{t-1}$, $\widehat{P}_t = D_t P_{t-1} D_t^{\top} + \Omega_t$ is the prediction step, $\kappa_t$ is the Kalman gain, and $(v_t - \widehat{S}_t^{\top} k_t)$ is the innovation. Note that the gain weights the residual in exactly the way the delta rule’s residual write is preserved. The Kalman gain is adjusted by evidence so that “when the memory is uncertain ( $\widehat{P}_t$ large), it writes strongly to new evidence, while protecting associations confirmed by repeated observations” (source: §3).

2) Existing Models Are Fixed-Gain Special Cases

Replacing the predicted covariance with an isotropic approximation $\widehat{P}_t \approx \widehat{b}_t I$ collapses the Kalman gain to the scalar $\beta_t k_t$ (source: §3.3):

In other words, DeltaNet ( $D_t = I$ ), Gated DeltaNet ( $D_t = \alpha_t I$ ), and KDA ( $D_t = \text{diag}(\alpha_t)$ ) differ only in their transition; they are fixed-gain Kalman filters that omit the covariance recursion and predict the gain directly from the token. This is the paper’s cleanest theoretical insight: forgetting is the “prediction” of the transition model, and write strength is the “confidence” of the covariance — two concepts cleanly separated.

3) Making It Practical: Two Scan-Compatible Approximations

The problem is that the exact Kalman update does not fit GPU-parallel linear attention. Each head would need to track a dense $d_k \times d_k$ covariance, and the gain follows a state-dependent Riccati recursion (source: §4). A parallel scan requires an associative update of the form $S_t = A_t S_{t-1} + b_t$.

To this end, the authors present two scan-compatible approximations.

• Diagonal KDN: restricts the covariance to a diagonal family and projects each token’s exact (dense) posterior onto a diagonal Gaussian via online mean-field variational inference (reverse KL minimization). This projection preserves the exact one-step posterior mean while maintaining $O(d_k)$ auxiliary state per head (source: Proposition 2, §4.1).
• Isotropic KDN: compresses uncertainty into a single scalar $b_t$ per head, using only $O(1)$ auxiliary state. The gain still takes the form $\kappa_t = \beta_t k_t$, but $\beta_t$ is now derived from the transition, process noise, and accumulated evidence (source: Proposition 4, §4.2).

The key point is that the uncertainty recursion is a Möbius map, so it can be expressed as a $2\times2$ matrix product and thus parallelized with a log-depth associative scan (source: §4.1):

How It Works: A Concrete Example

First, the full loop in one line:

  flowchart LR
  A["Current memory state S, P"] --> B["predict: transition D_t, noise Omega_t"]
  B --> C["predicted state S_pred, P_pred"]
  C --> D["observe value v_t with key k_t"]
  D --> E["innovation delta_t = v_t - predicted value"]
  E --> F["gain kappa_t = uncertainty / noise weighting"]
  F --> G["read o_t after update"]
  G --> A

Let’s grasp the intuition of the gain with a 2D toy example. Suppose the key space $d_k = 2$ and the current key points only at the first channel: $k_t = [1, 0]$. If the predicted uncertainty is $\widehat{p}_t = [p_1, p_2]$ (source: §4.1),

The gain acts only on channel 1, and its magnitude is set by the relative sizes of the uncertainty $p_1$ and the observation noise $r_t$. If $p_1 \gg r_t$ (this channel is still poorly known), then $\kappa_1 \to 1$, overwriting almost completely; if $p_1 \ll r_t$ (already confirmed), then $\kappa_1 \to 0$, ignoring the noisy observation. Existing KDA’s $\beta_t$ is predicted from the token in one shot without reflecting such evidence, so it cannot distinguish a “never-seen key” from a “key seen 100 times.”

The Möbius scan solves parallelization. The covariance update is expressed with the $2\times2$ matrix below, where composition becomes matrix multiplication (source: §4.1):

Unlike the sequential Riccati recursion of the exact Kalman filter, this associative operation is solved at log parallel depth with a prefix scan. The actual implementation is 3-pass: computing chunk-wise $2\times2$ summaries → exclusive scan for chunk entry states → chunk-parallel replay that emits gains. The memory update then becomes a fixed-gain input-only asymmetric delta rule, handled by a compact-WY kernel (source: §4.1, Appx D).

Information scale $\mu$. The diagonal approximation discards cross-channel correlations, so repeated keys can become overconfident in individual channels and cause excessive overwriting. The authors scale only the precision increase after a write by $\mu > 0$ to weaken future writes (source: §4.1). $\mu$ leaves the current write untouched (the gain is computed before writing) and protects only future writes. For normalized dense keys, $\mu = d_k$ compensates for the per-channel information dilution ($1/d_k$) (source: Appx B).

Performance: Key Results

All experiments pretrain from scratch on FineWeb-Edu, a controlled comparison that fixes the data, backbone, optimization recipe, and evaluation protocol, matching only model capacity (source: §5.1). Two scales are used: 750M/50B and 1.3B/100B.

Language modeling and commonsense reasoning (source: Tab. 1). At recurrent-only 750M, Diagonal KDN achieves WikiText PPL 18.64, LAMBADA PPL 14.15, and 6-task average accuracy 54.97%, beating both KDA (18.85/15.06/53.87) and Mamba-3 MIMO (18.99/15.67/54.39). At 1.3B, Diagonal KDN is again the best at 15.04/9.75/60.45%. Isotropic KDN delivers second-tier performance at every scale, notably improving over KDA while keeping almost the same cost structure.

In-context retrieval (source: Tab. 2, 3). On RULER needle-in-a-haystack, Diagonal KDN records the best aggregate among recurrent-only models at both scales. At 750M it maintains S-NIAH-1 at 100.0% up to 8K, and is notably strong on multi-key (MK-NIAH-1) and S-NIAH-3 (interference-inducing) — consistent with the design intent that per-channel uncertainty protects stored associations under interference (source: §5.2). On the 6-task real-world retrieval suite, Diagonal KDN also records the best recurrent-only average of 34.86% (leading on FDA, TriviaQA, and DROP).

Ablation (source: §5.3). The information scale peaks in average accuracy at $\mu = d_k$ (54.97%), but LAMBADA PPL is lowest at $\mu = 4d_k$ (13.91), so the effect is modest and metric-dependent. Fixing observation and process noise gave no consistent improvements. On efficiency, Isotropic KDN tracks KDA almost exactly, and Diagonal KDN stays close to GDN-2 while adding the per-channel uncertainty scan cost and maintaining linear scaling (source: Fig. 4). At long contexts, Mamba-3 SISO leads, and full attention degrades sharply as length grows.

Our Perspective: Strengths, Limitations, and Why It Matters

Strengths. The biggest contribution is conceptual unification. The scattered DeltaNet/GDN/KDA are folded into a single framework of “fixed-gain Kalman filters,” and the covariance (uncertainty) missing from them is promoted to a first-class citizen (source: §3.3). This is not mere renaming; it actually gives a principled answer (evidence-based gain) to the open question of “how to decide write strength.” The way the non-parallelizable exact Kalman filter is resolved with the elegant mathematical device of a Möbius scan is also impressive — it directly confronts the tension between a theoretically clean object (the Kalman filter) and hardware constraints (GPU parallel scans) (source: §4, Appx D).

Limitations. As the authors themselves admit, KDN is “not a complete realization of the Kalman associative memory, but a step in that direction” (source: §7). Exact filtering requires a dense covariance and a Riccati update, and the isotropic and diagonal approximations compress and discard information. Also, as the ablation shows, the information scale $\mu$ has a modest, metric-dependent effect, and the noise-fixing experiments show no clear gains, so the question “is the design principle the real cause of the performance gains” is not fully resolved (source: Tab. 4, 5). The experiment scale (up to 1.3B/100B) is small relative to current frontier LLMs, so whether the gains persist at scale is unknown. The architecture also borrows GDN-2’s backbone and is explicitly noted as not optimal (source: §5.1).

Why it matters. With long-context inference now a core cost of real services, fixed-size recurrent memory is one of the most promising alternatives to KV cache blowup. By introducing the quantitative principle of Bayesian/state estimation into that memory design, this paper gives a theoretical foundation to the empirically tuned delta-rule family. The insight that “uncertainty tracking is the write strategy” is a reusable framework for every architecture that uses memory compression, beyond recurrent attention.

What’s Next: The Road Ahead

The direct follow-up the authors propose is combining diagonal decay with Mamba-3’s damped rotation, allowing stored associations to be not only decayed but also rotated (source: §7). However, making richer transitions and covariance updates scan-efficient remains open.

Reasonable follow-ups naturally include (1) using the information form of the Kalman filter to approximate an associative scan over the dense covariance — going beyond Preconditioned DeltaNet’s static, fixed-noise limits to include learned stochastic dynamics (source: §6), (2) a channel-adaptive information scale that resolves $\mu$’s metric-dependent behavior, and (3) verifying whether the gains persist at scales of 7B and above. Ultimately, closing the remaining gap between “exact Kalman associative memory” and “parallel scans” is the most interesting open problem this work poses.

Tables from the paper

Tables converted mechanically from the arXiv e-print LaTeX source. The numbers are the paper’s own and did not pass through a model.

Table 1. Language modeling and zero-shot commonsense reasoning. Best per column within each block in bold, second best underlined; light-blue rows denote our KDN variants.

Table 2. In-context retrieval accuracy (%) on RULER single- and multi-key needle-in-a-haystack tasks . We use $25$-word context increments and deterministic random essay windows rather than $500$-word increments and fixed prefixes; every model receives the same $500$ samples per cell, with matched keys, values, and needle depths. Best per column within each block in bold, second best underlined.

Table 3. Zero-shot accuracy (%) on real-world retrieval tasks with inputs limited to $2$K tokens. Avg. is the unweighted six-task mean. Best per column within each block in bold, second best underlined.

Table 4. Mean effective write for the 750M Diagonal KDN ($d_k=128$). The first two columns use the matched random-initialization controls and checkpoints trained at each scale, evaluated on eight 2048-token sequences. The final column uses a separate 4096-token trajectory, changes only the runtime information scale of each trained checkpoint, and averages the resulting layer–head–token means across the four checkpoints; model weights are fixed, while downstream hidden states respond to the changed recurrence. Absolute levels should therefore be compared within, not across, these probe protocols. The intervention isolates the monotone protection effect from retraining compensation.

Figures in this post are taken from the original arXiv:2609.07816 (CC BY 4.0). Only size and format were changed.