Skip to content

Memory by Design: Probabilistic Sequence Layers

May 2026 · arXiv.org · Vol abs/2605.31163 · 0 citations · 45 references
Mathematics Computer Science

TL;DR

Distilling Bayesian Layers into a pretrained 340M Gated DeltaNet improves RULER long-context retrieval over a matched-compute control, and improves robustness beyond the training regime in controlled collision studies, learned associative recall, and the Zoology MQAR benchmark.

Abstract

We introduce the \emph{design-model framework}: a way to derive efficient recurrent sequence maps from explicit assumptions about memory. A design model writes evidence into memory by exact Bayesian filtering; a query- dependent readout produces a predictive distribution whose mean is the layer output. In our linear-Gaussian instantiation, the \emph{Bayesian Layer} propagates both a mean and a covariance: the covariance tracks uncertainty over stored associations, steering writes toward uncertain directions, attenuating gains as evidence accumulates, and preserving confident memories. The same framework unifies several sub-quadratic recurrences: linear attention, GLA, and Mamba-2/SSD are exact filters under a latent-input design model, whereas DeltaNet and related Delta-rule models are covariance-reset reductions of the Bayesian Layer's design model. Restoring covariance propagation yields closed-form predictions for retrieval dynamics, which we verify empirically, and improves robustness beyond the training regime in controlled collision studies, learned associative recall, and the Zoology MQAR benchmark. Training from scratch on WikiText-103 under matched state budgets lowers perplexity on associative-recall hits. Distilling Bayesian Layers into a pretrained 340M Gated DeltaNet improves RULER long-context retrieval over a matched-compute control, at a 2.5--2.7\% held-out perplexity cost.

View source

Similar papers

Preprint Aug 2026

Attention-Path Fragility as an Uncertainty Signal in Large Language Models

It is proposed that a model's uncertainty about a token is reflected not only in the breadth of its output distribution but also in whether a confident prediction is \emph{fragile} under perturbation of its attention pathways, a training-free estimator that masks attention heads and measures the BALD mutual information...

Minsoo Kim, Sungyoung Ji, Kisung Moon et al. · 0 citations
Preprint Aug 2026

Reading the Room: Implicit Confusion Encoding in Recurrent World Model States

This work confirms the signal is causally used, not merely present, by editing $h_t$ directly and watching behaviour change, including a check using real values from other trajectories instead of synthetic edits.

Donald Aadithiyan · 0 citations
Jul 2026

TransMem: Transforming Hidden States into Memory for Large Language Models

TransMem is proposed, a lightweight inference-time parametric memory module that transforms sparse historical hidden states from a frozen LLM backbone into reusable memory representations and introduces evidence-conditioned self-distillation to learn transferable memory utilization rather than task-specific knowledge.

Haodong Lei, Junming Liu, Yirong Chen et al. · 0 citations
#machine learning Preprint Aug 2026

Fast Weight Attention for Continual Learning

This framework separates temporal alignment, plasticity, forgetting, and bounded rehearsal in recurrent sequence models, together with numerically stable positive-decay renormalization, to remain competitive in language modeling and improve length extrapolation on variable-digit addition.

Yi-Fan Zhang, Steve Ta, Jasper Zhang et al. · 0 citations
Jul 2026

LatentFlow: A General Framework for Conditioning Stochastic Processes

LatentFlow is introduced, a single framework for conditioning stochastic processes, with no learned neural approximations and no training, that enables conditional sampling in seconds on a single desktop CPU across model classes that have never shared a scalable method.

Louis Sharrock, L. Astfalck, Henry B. Moss · 0 citations

Related blog posts

We use cookies to run the site and, with your consent, for analytics and to show ads. See our Cookie Policy.