PaperlayerView exact revision on arXiv

Source-linked AI summary

Kalman Delta Networks: Uncertainty-aware Associative Memory

Ngoc Bui, Tinglin Huang, Rex Ying

arXiv:2609.07816v1cs.LGcs.AI

TL;DR

Fixed-size linear-attention memories must overwrite associations without knowing future queries, while delta-rule gains do not track confidence in stored estimates. The paper reformulates this memory as a linear–Gaussian state-space model and introduces scan-compatible Kalman Delta Networks. Across two parameter scales, KDN variants improve perplexity and mean downstream accuracy over evaluated recurrent baselines.

  • Problem

    Fixed-size recurrent memories require online overwrite decisions, but delta-rule models do not adapt write strength to accumulated uncertainty in the memory estimate.

  • Method

    KDNs use Kalman filtering to propagate memory estimates and uncertainty, with Diagonal and Isotropic approximations designed for associative scans.

  • Results

    Both KDN variants achieve lower WikiText and LAMBADA perplexity and higher mean six-task zero-shot accuracy than every evaluated recurrent baseline at 750M/50B and 1.3B/100B.

  • Takeaways & Limitations

    Uncertainty-aware residual writes provide a scan-compatible extension of delta-rule recurrent memory, and Diagonal KDN attains the highest aggregate RULER score at both scales.

  • Takeaways & Limitations

    Exact uncertainty tracking requires dense covariance and a state-dependent Riccati recurrence, making it poorly suited to efficient parallel scans.

Abstract

from arXiv · show

Linear attention is increasingly used in frontier language models for efficient long-context inference and constant-memory decoding. Its fixed-size recurrent memory, however, requires an online decision at each token: what to write and how strongly to overwrite existing associations before knowing which information future queries will require. Delta-rule models learn this strength from the current token embedding but do not track confidence in the memory estimate, preventing each write from adapting to accumulated evidence. To represent this uncertainty explicitly, we reformulate recurrent associative memory as a linear--Gaussian state-space model, for which the Kalman filter is the optimal recursive estimator, and introduce a new family of models, Kalman Delta Networks (KDNs). Within KDNs, the transition propagates both the memory state and its uncertainty, allowing the Kalman gain to weight each residual write by accumulated evidence and observation reliability. Under this formulation, Delta-style updates emerge as a special case that substitutes a token-wise isotropic surrogate for predictive covariance and omits covariance tracking. Exact tracking, however, entails a dense, state-dependent Riccati recursion that is poorly suited to GPU-parallel linear-attention scans. To address this issue, we introduce two scan-compatible KDN approximations. Diagonal KDN projects each one-step posterior onto the diagonal Gaussian family through online mean-field variational inference, whereas Isotropic KDN uses an isotropic approximation with a single uncertainty scalar per head. Their uncertainty recurrences are Mobius maps, enabling associative scans with logarithmic parallel depth. Across controlled pretraining at 750M and 1.3B parameters, KDN variants consistently improve perplexity and mean downstream accuracy over state-of-the-art linear-attention models.

1. Introduction

Linear-attention memories replace quadratic history access with fixed-state online updates, but delta-rule gains lack explicit confidence. KDNs add uncertainty-aware Kalman filtering through scan-compatible approximations and improve evaluated language-model metrics.

  • Linear attention compresses the prefix into a fixed-size recurrent state, turning attention into an online memory-management problem.
  • Delta-rule mixers write prediction residuals, but their gains are predicted from the current token rather than accumulated memory confidence.
  • KDNs formulate recurrent memory as a linear–Gaussian state-space model whose Kalman update propagates both the memory estimate and covariance.
  • Exact Kalman filtering is impractical for linear attention because dense covariance tracking and state-dependent Riccati recurrences hinder parallel scans.
  • Diagonal and Isotropic KDN provide scan-compatible uncertainty approximations with O(dk) and O(1) auxiliary uncertainty state per head.
  • Both KDN variants improve perplexity and mean six-task zero-shot accuracy over evaluated recurrent mixers at 750M/50B and 1.3B/100B.

2. Preliminaries

Linear attention uses fixed-size associative memory to avoid growing caches, but additive updates cannot remove stale associations. Delta and gated variants improve editing or forgetting while retaining distinct limitations in selectivity and channel-wise retention.

  • Causal softmax attention retains preceding key–value pairs, whereas linear attention compresses history into a fixed-size recurrent state.
  • The additive linear-attention update can only write, so old associations cannot be explicitly removed and may interfere with new ones.
  • DeltaNet writes only the residual for the current key, but stale non-orthogonal associations can persist and interfere with later writes.
  • Gated DeltaNet applies data-dependent decay to the entire state, causing untouched associations to fade geometrically.
  • A single scalar decay cannot preserve long-lived information in some channels while rapidly forgetting transient information in others.
  • KDA gives each key channel its own decay rate, but its scalar erase/write gate still cannot use accumulated evidence to distinguish well-supported from uncertain associations.
  • Mamba adds control inputs directly, unlike delta mixers, which apply key-conditioned residual corrections.

3. Kalman Associative Memory

The paper formulates recurrent associative memory as a linear–Gaussian state-space model, making uncertainty an explicit state variable alongside the memory estimate. The resulting Kalman filter predicts memory evolution, incorporates noisy key–value observations through residual writes, and shows delta-rule updates as fixed-gain special cases.

  • Relation to delta-rule models: DeltaNet, Gated DeltaNet, and KDA are fixed-gain special cases that retain residual correction but omit covariance tracking and predict write strength from the current input.Using an isotropic predictive-covariance surrogate reduces the exact gain to the scalar β_t used by delta-rule models for normalized keys.
  • State-space formulation: The latent associative map evolves through a transition that models persistence, decay, and drift, while process noise represents unmodeled memory changes.For DeltaNet, Gated DeltaNet, and KDA, the transition is respectively I, α_tI, and diag(α_t).
  • State-space formulation: Each token provides a noisy value observation at one key, with observation covariance controlling how strongly the update should trust that value.Small r_t favors stronger writes; large r_t treats more of the value as context-specific noise.
  • Kalman filtering: Under linear–Gaussian assumptions, the Kalman filter is the optimal recursive estimator of the posterior mean and propagates both memory and key-space covariance.The filter solves the online MMSE problem by predicting the memory and then conditioning on the current key–value pair.
  • Kalman filtering: The Kalman update writes only the innovation—the observed value not explained by predicted memory—and uses covariance to determine the residual write’s direction and strength.High uncertainty permits stronger editing, whereas low uncertainty makes the prior memory resist noisy measurements.

4. Kalman Delta Networks

KDN approximations retain uncertainty-aware Kalman residual writes while making covariance tracking compatible with associative linear-attention scans. Diagonal KDN uses a diagonal variational state and Isotropic KDN uses a scalar uncertainty state, with uncertainty updates implemented through scan-friendly recurrences.

  • Motivation: Exact Kalman filtering is incompatible with fully parallel linear attention because covariance follows a Riccati recursion and requires dense dk × dk state per head.The gain depends on accumulated posterior uncertainty rather than input-only coefficients.
  • Diagonal KDN: Diagonal KDN restricts predictive covariance to a diagonal family and projects each dense posterior back using online mean-field variational inference.The projection preserves the exact one-step posterior mean conditional on the diagonal predictive prior.
  • Scan compatibility: Diagonal KDN reduces uncertainty tracking to dk scalars and replaces sequential Riccati updates with an associative scan of logarithmic parallel depth.Each channel evolves independently through a Möbius recurrence represented by 2 × 2 matrix composition.
  • Information scaling: Information scaling controls excessive overwrite caused by discarded cross-channel correlations and coordinatewise mean-field overconfidence.Increasing µ leaves the current write unchanged while reducing later effective writes; µ = 1 recovers the variational update.
  • Isotropic KDN: Isotropic KDN tracks one uncertainty scalar per head and derives its write strength from transition, process uncertainty, and accumulated evidence.This retains a scalar-gated update form while replacing the independent token gate with an uncertainty-derived strength.

5. Experiments

Experiments compare parameter-matched KDN mixers with recurrent and hybrid baselines across language modeling, reasoning, retrieval, ablations, and throughput. KDN variants generally improve quality while retaining linear scaling, with uncertainty representation affecting trade-offs.

  • Experimental setup: 750M/50B and 1.3B/100B parameter-matched pretraining fixes data, backbone, optimization, capacity, and evaluation protocols across models.Models are pretrained on FineWeb-Edu and evaluated on language modeling, commonsense reasoning, and retrieval tasks.
  • Language modeling and reasoning: Both KDN variants outperform KDA and Mamba-3 on WikiText, LAMBADA, and average accuracy at both scales, with Diagonal KDN leading average accuracy.In hybrid models, Diagonal KDN + SWA achieves the highest average accuracy, and every recurrent–attention hybrid beats the attention-only Transformer.
  • In-context retrieval: Diagonal KDN achieves the highest recurrent-only synthetic retrieval aggregate at both scales and the best recurrent-only average on six real-world retrieval tasks at 1.3B.Isotropic KDN + SWA has the best hybrid retrieval average, while every hybrid model substantially outperforms its recurrent-only counterpart.
  • Ablations: 4dk yields the best LAMBADA perplexity, while dk maximizes WikiText perplexity and average accuracy; the learned scale is metric-dependent.The information-scale effect is small overall, and no single setting dominates every metric.
  • Ablations: Neither fixing observation noise alone nor fixing both observation and process noise consistently improves short-context quality.Fixing observation noise improves WikiText perplexity but slightly worsens LAMBADA perplexity and average accuracy.
  • Throughput: Diagonal KDN retains linear throughput scaling despite extra channel-wise uncertainty-scan cost, while Isotropic KDN closely tracks KDA across sequence lengths.Mamba-3 SISO leads at long contexts, whereas full attention degrades sharply as sequence length grows.

6. Related Work

Related work spans fixed-memory attention, structured state-space sequence mixers, and Bayesian memory. KDN differs by treating token values as noisy observations and weighting residual updates with propagated uncertainty.

  • Fixed-memory attention: Fixed-memory alternatives reduce attention costs through sparse or bounded-cache attention, low-rank projection, and recurrent linear attention.Standard self-attention retains token-level history and computes quadratic pairwise interactions.
  • State-space mixers: Structured state-space models make long-range dynamics tractable through structured convolutions, scan-parallel recurrences, input selectivity, and attention–SSM duality.Mamba-3 extends this lineage with exponential–trapezoidal discretization, complex dynamics, and SISO/MIMO variants.
  • Bayesian memory: Unlike Mamba’s deterministic control-driven updates, KDN treats token values as noisy observations of a stochastic key–value map and weights innovations by propagated covariance.This connects KDN to Bayesian memory while preserving recurrent fixed-memory processing.

7. Conclusion and Future Work

The paper reframes delta-rule memory as Kalman innovation filtering and derives scan-compatible KDN approximations. Controlled experiments show broad gains, while exact Kalman memory remains computationally impractical and richer transitions remain future work.

  • Conclusion: KDN reformulates recurrent key–value memory as a linear–Gaussian dynamical system and derives gains from propagated uncertainty and observation noise.The delta-rule residual becomes the Kalman innovation: the observed value not explained by predicted memory.
  • Conclusion: Both Isotropic and Diagonal KDN achieve lower WikiText and LAMBADA perplexity and higher mean six-task zero-shot accuracy than evaluated recurrent baselines at 750M and 1.3B.Diagonal KDN also attains the highest aggregate RULER score at both scales.
  • Future work: Exact filtering maintains dense key-space covariance and a state-dependent Riccati update, so KDN approximations compress covariance to preserve parallel training.Future work includes transitions with damped rotations so stored associations can rotate as well as decay.

A.1. Proof of the Kalman optimal update

The proof derives the exact Kalman associative-memory update by predicting a Gaussian memory state, conditioning on the current key–value observation, and establishing its posterior and minimum-MSE gain.

  • State-space formulation: The latent associative-memory state follows a linear–Gaussian state-space model whose filtering posterior contains a mean estimate and key-space uncertainty.The posterior mean summarizes memory after previous key–value pairs, while covariance represents uncertainty.
  • Prediction: Before observing token t, the filter marginalizes the previous posterior to obtain a Gaussian predictive distribution.Linearity and Gaussianity make the prediction step analytically tractable.
  • Innovation: The current key reads a predicted value, and the difference between observed and predicted value forms the innovation used for updating memory.The observation consists of the current key–value pair.
  • Posterior update: Gaussian conditioning yields the posterior mean and covariance, with factorization across value coordinates under isotropic observation noise.Stacking the columnwise posteriors gives the associative-memory update.
  • Optimality: Completing the square shows that the Kalman gain uniquely minimizes the minimum-MSE objective over candidate gains.The exact update therefore has both a Bayesian conditioning and an optimization characterization.

A.2. Proof of the online diagonal variational update

The online diagonal variational update minimizes Gaussian KL over diagonal covariances, recovering the exact posterior mean while selecting diagonal variances through the factorized objective.

  • Proof: The Gaussian KL objective separates into a mean term and variance terms under the diagonal approximation.The determinant term is independent of the variational mean and variances.
  • Mean update: The variational mean is uniquely minimized at the exact posterior mean S = S⋆.Positive-definiteness of the exact posterior covariance makes the mean minimizer unique.
  • Variance update: Each diagonal variance is optimized independently against the corresponding exact posterior precision.The proof applies the scalar optimization separately for every pi > 0.

B. Overwrite and Protection–Write Decoupling

This section interprets effective write strength as overwrite behavior and shows that information scaling leaves the current write unchanged while protecting future writes along shared key channels.

  • Overwrite: Effective write strength measures the fraction of the current residual removed at the same key.Values near one correspond to an almost complete same-key rewrite, termed an overwrite.
  • Overwrite: Diagonal KDN at µ = 1 concentrates effective writes near one, while increasing µ progressively reduces write strength.This pattern is observed across layers, heads, and tokens in matched 750M KDA and Diagonal KDN runs.
  • Projection effect: For repeated keys, diagonal projection can retain more variance than the exact posterior, producing a stronger correction.The projected effective strength increases with retained variance and exceeds the exact value when m > 1.
  • Protection–write decoupling: The information scale µ leaves the current gain and memory estimate unchanged because the gain uses predictive state before the scaled uncertainty update.Its effect begins with future gains and weakens later writes along channels shared with the current key.
  • Empirical effect: The matched 750M sweep reduced pooled mean βeff from 0.886 at µ = 1 to 0.857 at 4dk before training, and from 0.943 to 0.805 after training.The result supports interpreting µ as future-write control rather than increased current observation strength.

C. Isotropic Kalman Delta Network

Isotropic KDN compresses uncertainty to one scalar per head while retaining Kalman residual writes, using variational projections and Möbius recurrences for scan-compatible computation.

  • Model definition: Isotropic KDN replaces the dk diagonal uncertainty values with one scalar uncertainty per head.The isotropic family approximates covariance as bI rather than retaining channelwise uncertainty values.
  • Projection structure: Isotropic projection is applied after prediction and after conditioning because the isotropic family is not closed under either Kalman stage.Channel-wise transitions can produce anisotropic predictive covariance, while conditioning produces dense posterior covariance.
  • Variational update: The isotropic variational update preserves the exact one-step posterior mean while approximating its covariance with btI.The retained covariance is obtained by minimizing Gaussian KL within the isotropic family.
  • Scan implementation: The scalar uncertainty recurrence is a Möbius map and therefore supports an associative prefix scan with parallel chunk replay.The implementation computes chunk summaries, scans exclusive carries, and replays chunks to emit gains.
  • Timing of uncertainty updates: The token-t information increment affects only subsequent gains because the gain is formed from pre-write predictive covariance.Reported Isotropic KDN runs fix the information scale to µ = dk.
  • Comparison with Diagonal KDN: Isotropic KDN preserves the predict–update decomposition and Kalman residual write but constrains the gain to κt = βtkt.Diagonal KDN instead retains a vector uncertainty state and permits an anisotropic gain direction.
Loading 2609.07816v1…