门控 KalmaNet：通过测试时岭回归实现的衰减记忆层
Gated KalmaNet: A Fading Memory Layer through Test-time Ridge Regression
GKA 用测试时岭回归加输入相关门控替代 softmax 注意力，在常数内存、线性复杂度下保留完整历史，可直接替换 ViT/混合架构中的注意力层；其 bfloat16 下的数值稳定化与分块切比雪夫迭代、自定义反向 kernel 对低精度长序列视觉骨干训练有直接复用价值。
原文依据 we show GKA outperforms Mamba when extended for ImageNet classification
摘要译文
由模型翻译 · 术语统一 · 请对照原文核对关键结论作为 softmax 注意力的高效替代方案，线性状态空间模型（SSM）具有常数内存和线性计算量，但仅维护对过去的有损、衰减式摘要，在面向回忆的设定中往往导致性能不佳。我们提出门控 KalmaNet（GKA），该层在预测下一个 token 时考虑完整历史，从而缩小这一差距，同时保持 SSM 式的效率。GKA 受卡尔曼滤波启发，在测试时在线求解岭回归问题，内存恒定且时间复杂度随序列长度线性增长。一个洞见是：标准卡尔曼滤波方程在低精度环境（如 bfloat16）中数值不稳定，且难以在现代 GPU 上并行化。我们通过两项创新同时应对这两个挑战：（1）一种带输入相关门控的自适应正则化策略，用于控制问题的条件数，从而确保数值稳定性并平衡记忆保留；（2）使用切比雪夫迭代替代传统迭代求解器，我们证明其在低精度设定下更为稳定。为提升可扩展性，我们以硬件感知的分块方式实现切比雪夫迭代，并针对自适应正则化与门控机制的反向传播编写了自定义 kernel。在短上下文任务上，GKA 展现出较强的语言理解能力，并优于现有 SSM（如 Mamba2 和 Gated DeltaNet）。在长上下文任务上，GKA 在高达 128k token 的真实 RAG 与 LongQA 任务中表现突出，相对基线提升超过 10%。最后，我们证明 GKA 在扩展用于 ImageNet 分类时优于 Mamba。
英文原文摘要 1673 字符
As efficient alternatives to softmax Attention, linear state space models (SSMs) achieve constant memory and linear compute, but maintain only a lossy, fading summary of the past, often leading to inferior performance in recall oriented settings. We propose Gated KalmaNet (GKA), a layer that reduces this gap by accounting for the full past when predicting the next token, while maintaining SSM-style efficiency. GKA is inspired by the Kalman Filter and solves ridge regression problems online at test time, with constant memory and linear time in the sequence length. An insight is that standard Kalman filter equations are numerically unstable in low-precision environments (e.g., bfloat16) and difficult to parallelize on modern GPUs. We address both challenges via two innovations: (1) an adaptive regularization strategy with input-dependent gating that controls the condition number of the problem, ensuring numerical stability and balancing memory retention; (2) the use of Chebyshev Iteration instead of conventional iterative solvers, which we show to be more stable in low-precision settings. To improve scalability, we implement Chebyshev Iteration in a hardware-aware, chunk-wise manner, along with custom kernels for backpropagating through our adaptive regularization and gating mechanisms. On short-context tasks, GKA shows strong language understanding capabilities and outperforms existing SSMs (e.g., Mamba2, and Gated DeltaNet). On long-context tasks, GKA excels at real-world RAG and LongQA tasks up to 128k tokens with more than 10% relative improvement over baselines. Finally, we show GKA outperforms Mamba when extended for ImageNet classification.