Memory by Design: Bayesian Layers for Sequence Modeling

To battle softmax attention's linearly growing inference memory cost, current recurrent architectures compress history into a fixed-dimensional state. Their recurrence equations alone, however, leave the underlying assumptions about memory implicit. We introduce design models: auxiliary probabilistic models specified by an initial distribution and input-dependent factors encoding memory dynamics and incoming evidence. Once these assumptions are specified, Bayesian filtering determines the recurrent update under the model. Specifically, a linear-Gaussian design model yields the Bayesian Layer, which maintains a memory estimate and its covariance via a Kalman update. Resetting this covariance recovers Delta-rule updates, while a different design model yields the additive recurrences of linear attention, GLA, and Mamba-2. These connections make explicit the assumptions underlying familiar memory updates. We illustrate how the geometry of the Bayesian Layer's uncertainty shapes subsequent updates and helps preserve earlier associations. In associative recall tasks, the Bayesian Layer improves recall on longer sequences and with greater overlap between keys than seen during training. We further evaluate its performance on multi-query associative recall (MQAR) and language modeling. Distilling Bayesian Layers into a pretrained Gated DeltaNet improves retrieval accuracy on RULER over a continued-training control.