Gated KalmaNet：通过测试时岭回归实现的衰减记忆层
摘要
线性状态空间模型（SSMs）为softmax注意力机制提供了一种高效的替代方案，具有恒定的内存和线性的计算量，但其对过去信息的有损、衰减式摘要损害了面向回忆的任务。我们提出了Gated KalmaNet（GKA，发音为“gee-ka”），这是一个能够考虑完整过去信息同时保持SSM风格效率的层。我们将该方法建立在卡尔曼滤波器（KF）的基础上，并证明了几种现有的SSM层（DeltaNet、Gated DeltaNet、Kimi Delta Attention）是在单位误差协方差假设下对KF递推的近似，该假设忽略了过去的键和值应如何最优地影响状态更新。相比之下，GKA维护完整的误差协方差并计算精确的卡尔曼增益。在能够实现并行化的稳态假设下，这简化为一个具有恒定内存和线性计算量的在线岭回归。标准的KF方程在低精度设置（例如bfloat16）下数值不稳定，且难以在GPU上并行化。我们通过以下方式解决此问题：（1）通过输入依赖的门控进行自适应正则化，以控制岭回归的条件数；（2）切比雪夫迭代，我们证明其在低精度下比传统的迭代求解器更稳定。我们进一步开发了硬件感知的分块内核以实现高效训练。实验证明，GKA在短上下文任务上优于现有的SSM层（例如Mamba2、Gated DeltaNet），并在长达128k tokens的长上下文RAG和LongQA任务上实现了超过10%的相对改进。我们进一步证明，当扩展到ImageNet分类时，GKA优于Mamba。我们的代码，包括用于训练和推理（vLLM）的Triton内核，以及HuggingFace上基于GKA的8B和32B规模混合模型库，均以Apache 2.0许可证发布。
引用
@article{arxiv.2511.21016,
title = {Gated KalmaNet: A Fading Memory Layer Through Test-Time Ridge Regression},
author = {Liangzu Peng and Aditya Chattopadhyay and Luca Zancato and Elvis Nunez and Wei Xia and Stefano Soatto},
journal= {arXiv preprint arXiv:2511.21016},
year = {2026}
}
备注
30 pages, 10 figures. Accepted at CVPR 2026