# Kalman Delta Networks: Uncertainty-aware Associative Memory

> Kalman Delta Networks reformulate delta-rule associative memory as a Kalman filter that tracks uncertainty, outperforming linear-attention baselines in perplexity and retrieval.

- **Source:** [arXiv](https://arxiv.org/abs/2609.07816)
- **Published:** 2026-09-12
- **Permalink:** https://picx.dev/p/lalIEg
- **Whiteboard:** https://picx.dev/p/lalIEg/image

## Summary

# Kalman Delta Networks: Uncertainty-aware Associative Memory

## Summary (Overview)

- The paper introduces **Kalman Delta Networks (KDNs)**, a new family of linear-attention models that reformulate recurrent associative memory as a linear–Gaussian state-space model, for which the Kalman filter is the optimal recursive estimator.
- KDNs explicitly track **uncertainty (covariance)** alongside the memory state, enabling each write to be weighted by accumulated evidence and observation reliability—a capability missing from existing delta-rule models (DeltaNet, Gated DeltaNet, KDA).
- Two scan-compatible approximations are derived: **Diagonal KDN** (one uncertainty value per key channel, $O(d_k)$ auxiliary state per head) via online mean-field variational inference, and **Isotropic KDN** (single uncertainty scalar per head).
- An **information-scaling factor** ($\mu$) is introduced to mitigate excessive overwrite caused by the diagonal approximation's underprotection of stored key directions.
- Controlled pretraining experiments at 750M and 1.3B parameters show consistent improvements in perplexity and mean downstream accuracy over state-of-the-art linear-attention baselines, including Mamba-3 and DeltaNet variants.

---

## Introduction and Theoretical Foundation

### Background: Linear Attention and Fixed-Size Memory

Self-attention allows every query to retrieve values associated with all preceding keys but requires a cache that grows linearly with sequence length and quadratic query–key interactions. Linear attention addresses this by compressing token history into a **fixed-size recurrent state** $S_t \in \mathbb{R}^{d_k \times d_v}$:

$$S_t = S_{t-1} + v_t k_t^\top$$

Each token writes a key–value association to the state, and each query reads from it. However, additive updates can only write—old associations cannot be explicitly removed and may interfere with new ones.

### Delta-Rule Models and Their Limitations

Delta-rule models address this by writing only the **residual**—the part of the value not explained by the current memory:

$$S_t = S_{t-1} + \beta_t (v_t - S_{t-1}^\top k_t) k_t^\top$$

where $\beta_t$ is a scalar write strength. Key limitations of existing delta-rule models:

- **DeltaNet**: Uses scalar gate $\beta_t$ predicted from the current token; no decay of stale associations.
- **Gated DeltaNet**: Adds scalar decay $\alpha_t$ applied uniformly to all key channels.
- **KDA (Key-Value Delta Attention)**: Uses diagonal transition $D_t = \mathrm{diag}(\boldsymbol{\alpha}_t)$ for channel-wise decay but still controls write strength with a scalar gate $\beta_t$ predicted from the current token.

> **Critical limitation**: The write strength in all delta-rule models is predicted from the current token embedding rather than derived from **explicit confidence in stored associations**. These models cannot distinguish a well-supported association from an uncertain one.

### Theoretical Reformulation: Linear-Gaussian State-Space Model

The paper casts recurrent associative memory as a **linear–Gaussian state-space model** where the latent memory $\widetilde{S}_t$ evolves as:

$$\widetilde{S}_t = D_t \widetilde{S}_{t-1} + \Omega_t^{1/2} \varepsilon_t \tag{9}$$

with observation model:

$$v_t = \widetilde{S}_t^\top k_t + r_t^{1/2} \epsilon_t \tag{10}$$

where:
- $D_t$ is the transition (persistence/decay of associations)
- $\Omega_t$ is the process noise covariance (how much the memory drifts)
- $r_t$ is the observation noise (reliability of the current token's value)
- $k_t \in \mathbb{R}^{d_k}$ is the key, $v_t \in \mathbb{R}^{d_v}$ is the value

Under linear–Gaussian assumptions, the **Kalman filter** is the optimal recursive estimator. The posterior mean update takes the form:

$$\hat{S}_t = D_t S_{t-1} + \kappa_t (v_t - k_t^\top D_t S_{t-1}) k_t^\top \tag{17}$$

where $\kappa_t$ is the Kalman gain:

$$\kappa_t = \frac{P_t^- k_t}{k_t^\top P_t^- k_t + r_t}$$

This is a **residual write**—the filter adds only the part of the observed value that the predicted memory failed to explain—but with the crucial difference that the write strength is derived from the **predictive covariance** $P_t^-$, which tracks accumulated evidence.

### Connection to Existing Models

Delta-rule models are identified as **fixed-gain special cases** of the Kalman filter. For normalized keys, the exact Kalman gain collapses to the scalar write strength $\beta_t$ used by delta-rule models. The models differ only in how the memory is predicted before the write:

| Model | Transition $D_t$ | Gain Source |
|-------|-----------------|-------------|
| DeltaNet | $I$ (identity) | Token-predicted scalar |
| Gated DeltaNet | $\alpha_t I$ (scalar decay) | Token-predicted scalar |
| KDA | $\mathrm{diag}(\boldsymbol{\alpha}_t)$ | Token-predicted scalar |
| **KDN** | **Learned diagonal** | **Covariance-derived (Kalman gain)** |

---

## Methodology

### Exact Kalman Filter for Associative Memory

The exact Kalman update tracks both the mean $S_t$ and covariance $P_t$ of the latent memory. The filter proceeds in two steps:

**Prediction step** (before observing token $t$):

$$\widetilde{S}_t = D_t S_{t-1}, \quad P_t^- = D_t P_{t-1} D_t^\top + \Omega_t$$

**Update step** (conditioning on observation $(k_t, v_t)$): Using Gaussian conditioning on each value coordinate jointly with the scalar observation $v_{t,j}$:

$$S_t = \widetilde{S}_t + \kappa_t (v_t - k_t^\top \widetilde{S}_t) k_t^\top \tag{17}$$

$$P_t = P_t^- - \kappa_t P_t^- k_t k_t^\top \tag{18}$$

$$\kappa_t = \frac{P_t^- k_t}{k_t^\top P_t^- k_t + r_t} \tag{19}$$

**Problem**: The exact recursion is not scan-compatible because:
1. $P_t$ follows a **Riccati recursion** (non-linear)
2. The gain $\kappa_t$ depends on accumulated posterior uncertainty
3. Tracking $P_t$ requires a **dense $d_k \times d_k$ state** per head

### Diagonal KDN: Online Mean-Field Variational Inference

The first approximation restricts the tracked covariance to the **diagonal family**. After each token, the posterior is projected onto the diagonal Gaussian family via online mean-field variational inference:

$$P_t \approx \mathrm{diag}(p_t)$$

This yields a per-channel uncertainty vector $p_t \in \mathbb{R}^{d_k}$ with $O(d_k)$ auxiliary state per head.

**Information scaling**: To mitigate excessive overwrite under the diagonal approximation (which can underprotect stored key directions), an information scale $\mu > 0$ is applied only to the post-write precision increment:

$$p_t^{-1} = (D_t \odot p_{t-1})^{-1} + \mu \cdot \frac{k_t \odot k_t}{r_t}$$

where $\odot$ denotes element-wise operations. In experiments, $\mu = d_k$ is used as the fixed information scale.

### Isotropic KDN: Single Scalar Uncertainty

The second approximation restricts the covariance to the **isotropic family**:

$$P_t \approx b_t I$$

Unlike the diagonal family, the isotropic family is **not closed** under a channel-wise transition, so Isotropic KDN projects both the predicted covariance and the measurement posterior back to this family. This retains the scalar-gated form used by DeltaNet/Gated DeltaNet/KDA:

$$\beta_t = \frac{\widehat{b}_t}{b_t + r_t}$$

but now $\beta_t$ depends on the transition, process uncertainty, and evidence accumulated from previous tokens through $\widehat{b}_t$. This requires only $O(1)$ auxiliary state per head.

### Training and Evaluation Setup

- **Scales**: 750M and 1.3B parameters
- **Data**: Controlled pretraining with matching data, backbone, optimization recipe, and evaluation protocol
- **Capacity matching**: Recurrent-only models use the same backbone with feed-forward widths adjusted to match non-embedding parameter counts within 0.03%; Mamba-3 baselines match KDN total recurrent-state size
- **Sequence length**: 4K tokens, global batch of 0.5M tokens
- **Evaluation**: Perplexity, downstream tasks, and RULER in-context retrieval benchmarks

---

## Empirical Validation / Results

### Language Modeling Perplexity

KDN variants consistently improve perplexity over state-of-the-art linear-attention models at both 750M and 1.3B scales, including DeltaNet, Gated DeltaNet, KDA, and Mamba-3 baselines.

### In-Context Retrieval (RULER)

**Table 2** presents in-context retrieval accuracy (%) on RULER single- and multi-key needle-in-a-haystack tasks (25-word context increments with deterministic random windows):

| Model | Single-Key | Multi-Key |
|-------|-----------|-----------|
| Best delta-rule baseline | (second best) | (second best) |
| **Diagonal KDN** | **Best** | **Best** |
| **Isotropic KDN** | **Best/second** | **Best/second** |

KDN variants achieve the best or second-best performance per column across all retrieval tasks, demonstrating that uncertainty-aware writes improve the ability to store and retrieve associations in long contexts.

### Ablations: Information Scale

**Table 4** ablates the information scale using fixed values $\mu \in \{1, \sqrt{d_k}, d_k, 4d_k\}$ and a learned scale initialized at $d_k$. Results show that:
- $\mu = d_k$ performs well across settings
- The learned scale initialized at $d_k$ provides competitive or better performance
- Too small $\mu$ leads to excessive overwrite; too large $\mu$ makes the model underwrite new information

---

## Theoretical and Practical Implications

### Theoretical Contributions

1. **Unified view of delta-rule and state-space models**: The paper bridges delta-rule recurrent mixers (DeltaNet, Gated DeltaNet, KDA) and state-space models (Mamba lineage) by showing delta-rule updates are fixed-gain special cases of Kalman filtering over a latent key–value map. Unlike Mamba's control-driven additive input, these models correct a predicted key–value map with a key-conditioned residual.

2. **Uncertainty as the missing state variable**: The key insight is that delta-rule models omit covariance tracking, which is essential for principled write-strength determination. KDNs derive the innovation gain jointly from uncertainty and the transition, rather than from a token-predicted scalar.

3. **Scan-compatible uncertainty tracking**: The paper demonstrates that uncertainty-aware filtering can be made compatible with GPU-parallel linear-attention scans through diagonal (mean-field) and isotropic approximations, with $O(d_k)$ and $O(1)$ auxiliary state per head respectively.

### Practical Implications

- **Constant-memory decoding with uncertainty awareness**: KDNs maintain the efficient constant-memory decoding and scan-parallel training of linear attention while adding explicit confidence tracking.
- **Improved retrieval and modeling**: The empirical gains on RULER tasks and perplexity suggest that uncertainty-aware writes better balance the trade-off between preserving old associations and incorporating new evidence.
- **Information scaling as a practical calibration tool**: The $\mu$ parameter provides a simple, interpretable knob for controlling the overwrite behavior of diagonal approximations.

### Relationship to Prior Work

- **Kalman Linear Attention (KLA)** also tracks uncertainty with a scan-parallel information-form filter but factorizes its belief over feature/state-expansion coordinates under time-invariant OU dynamics. KDNs instead maintain exact posterior covariance in key space (dense, shared across value channels) with isotropic/diagonal scan-compatible approximations.

---

## Conclusion

This work bridges delta-rule recurrent mixers and state-space models by reformulating recurrent key–value memory as a **linear–Gaussian dynamical system over a latent, non-stationary key–value map**. The delta-rule residual write is given a principled interpretation as the **Kalman innovation**—the part of the observed value not explained by the predicted memory.

### Key Takeaways

1. Delta-rule models are fixed-gain Kalman filters that omit covariance tracking; uncertainty is their missing state variable.
2. Scan-compatible Isotropic and Diagonal KDNs make uncertainty-aware filtering practical for linear attention.
3. Information scaling limits future overwrite of stored associations under the diagonal approximation.
4. KDN variants consistently outperform state-of-the-art linear-attention baselines in perplexity and downstream accuracy.

### Future Directions

- **Richer transitions**: Extending diagonal decay with damped rotations (as in Mamba-3) to allow stored associations to rotate as well as decay.
- **Scan-efficient richer covariance**: Making the covariance update for richer transition models scan-efficient remains open.
- The authors note that KDNs are "steps toward, rather than full realizations of, Kalman Associative Memory"—full covariance tracking with scan-compatible approximations remains an open challenge.

---

_Markdown view of https://picx.dev/p/lalIEg, served by PicX — AI-generated visual whiteboard summaries of research papers._