让记忆网络学会评估自身不确定性,提升长文本推理效果。
Kalman Delta Networks: Uncertainty-aware Associative Memory

- 将记忆更新建模为卡尔曼滤波,显式跟踪记忆的不确定性。
- 在7.5亿和13亿参数模型上,比现有线性注意力模型更优。
- 适合需要可靠记忆推理的长序列任务,如对话、文档理解。
线性注意力广泛用于前沿语言模型,以实现高效长上下文推理和恒定内存解码。然而,其固定大小的循环记忆需在每个词元处做出在线决策:写入什么内容以及写入强度,但此时尚不知未来查询所需信息。增量规则模型从当前词元嵌入中学习写入强度,却不追踪记忆估计的置信度,导致每次写入无法根据累积证据调整。为此,我们把循环关联记忆重新建模为线性-高斯状态空间模型,其中卡尔曼滤波是最佳递归估计器,并提出新模型族——卡尔曼增量网络(KDN)。在KDN中,转移过程同时传播记忆状态及其不确定性,使卡尔曼增益可根据累积证据与观测可靠性加权每次残差写入。在此框架下,传统的增量更新仅为一种特殊情况,即用词元级各向同性近似替代预测协方差并省略协方差追踪。精确追踪则涉及密集的、依赖状态的里卡蒂递推,不适合GPU并行扫描。为解决此问题,我们引入两种扫描兼容的近似方法:对角线KDN通过在线均值场变分推断将每一步后验投影至对角高斯族;各向同性KDN采用单个不确定性标量表示每头的不确定度。它们的不确定性递推为莫比乌斯映射,支持对数级并行深度的关联扫描。在7.5亿和1.3亿参数的受控预训练中,各类KDN持续优于最先进线性注意力模型,在困惑度和下游平均准确率上均有提升。
原文摘要 · Abstract (English)
Linear attention is increasingly used in frontier language models for efficient long-context inference and constant-memory decoding. Its fixed-size recurrent memory, however, requires an online decision at each token: what to write and how strongly to overwrite existing associations before knowing which information future queries will require. Delta-rule models learn this strength from the current token embedding but do not track confidence in the memory estimate, preventing each write from adapting to accumulated evidence. To represent this uncertainty explicitly, we reformulate recurrent associative memory as a linear--Gaussian state-space model, for which the Kalman filter is the optimal recursive estimator, and introduce a new family of models, Kalman Delta Networks (KDNs). Within KDNs, the transition propagates both the memory state and its uncertainty, allowing the Kalman gain to weight each residual write by accumulated evidence and observation reliability. Under this formulation, Delta-style updates emerge as a special case that substitutes a token-wise isotropic surrogate for predictive covariance and omits covariance tracking. Exact tracking, however, entails a dense, state-dependent Riccati recursion that is poorly suited to GPU-parallel linear-attention scans. To address this issue, we introduce two scan-compatible KDN approximations. Diagonal KDN projects each one-step posterior onto the diagonal Gaussian family through online mean-field variational inference, whereas Isotropic KDN uses an isotropic approximation with a single uncertainty scalar per head. Their uncertainty recurrences are Mobius maps, enabling associative scans with logarithmic parallel depth. Across controlled pretraining at 750M and 1.3B parameters, KDN variants consistently improve perplexity and mean downstream accuracy over state-of-the-art linear-attention models.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。