Memory by Design: Probabilistic Sequence Layers

📄 arXiv: 2605.31163v1 📥 PDF

作者: Matthew Dowling, Hyungju Jeon, Cristina Savin, Il Memming Park

分类: stat.ML, cs.LG

发布日期: 2026-05-29

备注: Preprint, in submission

💡 一句话要点

提出设计模型框架，通过显式记忆假设推导高效循环序列映射。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 循环序列模型 贝叶斯滤波 记忆网络 长上下文检索 不确定性建模

📋 核心要点

  1. 现有循环序列模型缺乏对记忆机制的显式建模，限制了其效率和可解释性。
  2. 论文提出设计模型框架，通过贝叶斯滤波将信息写入记忆，并利用协方差跟踪不确定性，指导记忆更新。
  3. 实验表明，该方法在长上下文检索和关联回忆等任务中表现出更强的鲁棒性，并能提升现有模型的性能。

📝 摘要（中文）

本文提出了一种设计模型框架，该框架通过对记忆的显式假设来推导高效的循环序列映射。设计模型通过精确的贝叶斯滤波将证据写入记忆；一个依赖于查询的读出操作产生一个预测分布，其均值作为层的输出。在线性-高斯实例化中，贝叶斯层传播均值和协方差：协方差跟踪存储关联的不确定性，引导写入朝向不确定的方向，随着证据的积累衰减增益，并保留置信的记忆。该框架统一了几个亚二次复杂度循环。线性注意力、GLA和Mamba-2/SSD是同一设计模型下的精确滤波器，而DeltaNet和相关的Delta-rule模型则是在另一个模型下协方差重置的简化版本。恢复协方差可以得到检索动态的闭式预测（经验证），并在受控的冲突研究、学习的关联回忆和动物学MQAR基准测试中提高了超出训练范围的鲁棒性；将贝叶斯层提炼到预训练的3.4亿参数的Gated DeltaNet中，提高了在匹配计算下的RULER长上下文检索性能。

🔬 方法详解

问题定义：现有循环序列模型在处理长序列时面临效率和鲁棒性挑战。传统的RNN及其变体难以捕捉长期依赖关系，并且容易受到梯度消失或爆炸的影响。注意力机制虽然能够缓解这些问题，但计算复杂度较高。此外，现有模型缺乏对记忆过程的显式建模，导致其行为难以解释和控制。

核心思路：论文的核心思路是将循环序列建模视为一个贝叶斯滤波过程，通过显式地对记忆进行建模，并利用贝叶斯推理来更新和检索记忆。这种方法能够有效地跟踪记忆中的不确定性，并根据不确定性来指导记忆的更新，从而提高模型的效率和鲁棒性。

技术框架：该方法的核心是设计模型框架，它包含两个主要模块：写入模块和读出模块。写入模块通过贝叶斯滤波将输入序列的信息写入记忆，并更新记忆的状态和不确定性。读出模块根据查询从记忆中检索相关信息，并生成预测分布。在线性-高斯实例化中，记忆的状态由均值和协方差表示，协方差用于跟踪记忆的不确定性。

关键创新：该方法最重要的创新点在于将循环序列建模与贝叶斯滤波联系起来，并利用协方差来跟踪记忆的不确定性。这种方法能够显式地建模记忆过程，并根据不确定性来指导记忆的更新，从而提高模型的效率和鲁棒性。此外，该框架还统一了多种亚二次复杂度循环模型，如线性注意力、GLA和Mamba-2/SSD。

关键设计：在线性-高斯实例化中，记忆的状态由均值和协方差表示。写入模块使用卡尔曼滤波来更新均值和协方差。读出模块根据查询从记忆中检索相关信息，并生成预测分布。损失函数可以使用负对数似然损失或交叉熵损失。关键参数包括记忆的大小、滤波器的参数和读出模块的参数。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，该方法在长上下文检索、关联回忆和动物学MQAR基准测试中表现出更强的鲁棒性。例如，将贝叶斯层提炼到预训练的3.4亿参数的Gated DeltaNet中，提高了在匹配计算下的RULER长上下文检索性能。此外，该方法还能够预测检索动态，并验证了其理论分析的有效性。

🎯 应用场景

该研究成果可应用于各种需要处理长序列数据的场景，如自然语言处理、语音识别、视频分析和时间序列预测。例如，可以用于改进机器翻译、文本摘要、对话系统和推荐系统等任务的性能。此外，该方法还可以用于开发更高效和鲁棒的循环神经网络，从而推动人工智能技术的发展。

📄 摘要（原文）

We introduce the 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 one design model, whereas DeltaNet and related Delta-rule models arise as covariance-reset reductions under another. Restoring the covariance yields closed-form predictions for retrieval dynamics, verified empirically, and improves robustness beyond the training regime across controlled collision studies, learned associative recall, and the Zoology MQAR benchmark; distilling Bayesian Layers into a pretrained 340M Gated DeltaNet improves RULER long-context retrieval at matched compute.