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:

---

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):

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.