Kalman Delta Networks: Uncertainty-aware Associative Memory

Kalman Delta Networks. Uncertainty-aware Associative Memory. By Ngoc Bui and two others, from Yale University. Published on the seventh of September. Its code repository is public and real: about twenty-five source files.

Linear attention is increasingly used in frontier language models, because it makes long-context inference efficient and keeps decoding memory constant. The price is a recurrent memory of fixed size, and that forces an online decision at every token: what to write, and how strongly to overwrite the associations already stored, all before knowing what future queries will need. Delta-rule models learn how strongly to write from the current token's embedding. What they do not track is confidence in the memory's own estimate. Without it, a write cannot adapt to how much evidence has already accumulated behind what the memory holds. A memory slot that has been confirmed many times and one written once are treated alike.

To represent that uncertainty explicitly, the authors recast recurrent associative memory as a linear-Gaussian state-space model. For such a model, the Kalman filter is the optimal recursive estimator. Building on that, they introduce a new family, Kalman Delta Networks, or KDNs. In a KDN, the transition carries forward both the memory state and its uncertainty. The Kalman gain can then weight each residual write by the evidence accumulated so far and by how reliable the observation is. A well-established association can resist being overwritten by a single token, while an uncertain one can be updated more freely, and a noisy observation counts for less than a reliable one.

The formulation also explains where Delta-style updates come from. They emerge as a special case: one that substitutes a token-wise isotropic surrogate for the predictive covariance, and that does not track covariance at all. The delta rule is what the Kalman view looks like with its uncertainty bookkeeping switched off.

Exact tracking has a cost. It requires a dense, state-dependent Riccati recursion, which is poorly suited to the GPU-parallel scans that linear attention relies on. The authors therefore introduce two approximations that fit those scans. Diagonal KDN projects each one-step posterior onto the diagonal Gaussian family, using online mean-field variational inference. Isotropic KDN uses an isotropic approximation, keeping a single uncertainty scalar per head. The first keeps a separate uncertainty for each dimension, the second collapses it to one number, trading detail for simplicity.

In both cases the uncertainty recurrence takes the form of Mobius maps. That property is what lets them run as associative scans with logarithmic parallel depth, so the added uncertainty tracking stays compatible with the parallel computation linear attention depends on.

In controlled pretraining at two scales, seven hundred and fifty million parameters and one point three billion parameters, KDN variants consistently improve both perplexity and mean downstream accuracy over state-of-the-art linear-attention models.

A memory write looks different after this. Instead of a learned strength applied blindly, it becomes an estimate update that knows how confident it is, and the familiar delta rule turns out to be the version that forgot to ask. The approximations are what make that principled update cheap enough to scan in parallel.