用卡尔曼滤波改进状态空间模型，让记忆更持久且高效。
Gated KalmaNet: A Fading Memory Layer Through Test-Time Ridge Regression
- 基于卡尔曼滤波设计新层，精确计算历史信息影响。
- 长文本任务上相对提升超10%，最高支持128k tokens。
- 适合需要长程记忆的推理与图像分类任务。
线性状态空间模型（SSMs）以常数内存和线性计算效率替代softmax注意力，但其对历史信息的损失性压缩不利于召回型任务。本文提出Gated KalmaNet（GKA），在保持SSM效率的同时完整保留历史信息。我们基于卡尔曼滤波（KF）构建方法，揭示现有SSM层（如DeltaNet、Gated DeltaNet、Kimi Delta Attention）在单位误差协方差假设下为KF的近似，忽略历史键值对状态更新的最优影响。GKA则维持完整误差协方差，精确计算卡尔曼增益。在稳态假设下可并行化，等价于在线岭回归，具备常数内存与线性计算复杂度。针对低精度（如bfloat16）下传统KF数值不稳定的难题，我们引入输入依赖门控实现自适应正则化，并采用切比雪夫迭代，相比常规求解器更稳定。此外开发了面向硬件的分块内核以提升训练效率。实验表明，GKA在短上下文任务中超越Mamba2、Gated DeltaNet；在长上下文RAG与LongQA任务上相对提升超过10%（最长达128k tokens）。扩展至ImageNet分类时亦优于Mamba。代码已开源，包含Triton训练/推理内核（vLLM支持）及HuggingFace上的8B与32B规模混合模型库，许可为Apache 2.0。
原文摘要 · Abstract (English)
Linear State-Space Models (SSMs) offer an efficient alternative to softmax Attention with constant memory and linear compute, but their lossy, fading summary of the past hurts recall-oriented tasks. We propose Gated KalmaNet (GKA, pronounced "gee-ka"), a layer that accounts for the full past while retaining SSM-style efficiency. We ground our approach in the Kalman Filter (KF), and show that several existing SSM layers (DeltaNet, Gated DeltaNet, Kimi Delta Attention) are approximations to the KF recurrence under an identity error covariance assumption, which ignores how past keys and values should optimally influence state updates. In contrast, GKA maintains the full error covariance and computes the exact Kalman gain. Under a steady-state assumption that enables parallelization, this reduces to an online ridge regression with constant memory and linear compute. The standard KF equations are numerically unstable in low-precision settings (e.g., bfloat16) and hard to parallelize on GPUs. We address this with (1) adaptive regularization via input-dependent gating to control the ridge regression's condition number, and (2) Chebyshev Iteration, which we show is more stable than conventional iterative solvers in low precision. We further develop hardware-aware chunk-wise kernels for efficient training. Empirically, GKA outperforms existing SSM layers (e.g., Mamba2, Gated DeltaNet) on short-context tasks and achieves more than 10\% relative improvement on long-context RAG and LongQA up to 128k tokens. We further show GKA outperforms Mamba when extended to ImageNet classification. Our code, including Triton kernels for training and inference (vLLM), along with a model zoo of GKA-based Hybrid models at 8B and 32B scale on HuggingFace, is released under Apache 2.0.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。