用概率设计框架构建高效序列模型,提升记忆精度与鲁棒性。
Memory by Design: Probabilistic Sequence Layers

- 基于贝叶斯滤波设计记忆机制,显式建模不确定性。
- 在控制碰撞和联想回忆任务中显著提升泛化能力,降低困惑度。
- 适用于长序列建模,适合追求高鲁棒性的大模型优化场景。
我们提出一种设计-模型框架,从对记忆的显式假设出发,推导出高效的递归序列映射。设计模型通过精确贝叶斯滤波将证据写入记忆,查询依赖的读取生成预测分布,其均值即为层输出。在线性高斯实例中,贝叶斯层同时传播均值与协方差:协方差追踪存储关联的不确定性,引导写入至不确定方向,随证据累积衰减增益,保留高置信记忆。同一框架统一了多个亚二次递归结构:线性注意力、GLA、Mamba-2/SSD 是潜在输入设计模型下的精确滤波器;而DeltaNet及类似增量规则模型则是贝叶斯层设计模型的协方差重置简化版。恢复协方差传播可得检索动态的闭式预测,经实证验证,并在受控碰撞实验、学习到的联想回忆及Zoology MQAR基准上表现出更强鲁棒性。在匹配状态预算下,于WikiText-103上从头训练,降低了联想回忆命中时的困惑度。将贝叶斯层蒸馏为预训练的340M门控DeltaNet,在相同计算量下优于对照组,长上下文检索性能提升,仅付出2.5–2.7%的持有率困惑度代价。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。