Gated KalmaNet: A Fading Memory Layer Through Test-Time Ridge Regression

AI summary

Overview

Research area: Efficient sequence modeling for large language models — specifically linear State-Space Model (SSM) layers that replace softmax Attention.

Technical level: Advanced. The paper assumes familiarity with state-space recurrences, ridge regression, the Kalman Filter, iterative linear solvers, condition numbers, and low-precision (bfloat16) training.

Scope: The paper proposes Gated KalmaNet (GKA), a layer that computes its state from the entire past while retaining the constant-memory and linear-compute profile of linear SSMs, and evaluates it on synthetic recall, short-context, long-context, and image classification tasks.

What This Paper Is About

Linear SSMs such as Mamba2 and DeltaNet are cheap to run — constant memory, linear compute — but they compress the past into a fixed-size "fading, lossy" state, which hurts tasks that require recalling specific earlier tokens. Softmax Attention avoids this loss by keeping the full KV-cache, but pays quadratic time and linearly growing storage. The paper's goal is a memory layer that behaves like Attention (conditioned on the exact past) at the cost of an SSM, by grounding the layer in the Kalman Filter and solving its underlying test-time ridge regression stably and in parallel.

Key Contributions

1. A Kalman-Filter interpretation of existing SSMs. The authors show that DeltaNet, Gated DeltaNet, and Kimi Delta Attention can be viewed as approximations to the Kalman Filter recurrence under an identity error-covariance assumption — which discards information about how past keys and values should optimally influence state updates. GKA instead maintains the full error covariance and computes the exact Kalman gain, reducing under a steady-state assumption to an online ridge regression with constant memory and linear compute.

2. Adaptive regularization plus Chebyshev Iteration (CH) for numerical stability. The regularization strength is set adaptively as λ_t = a · ||H_t||_F, which upper-bounds the condition number at (a+1)/a. An exact solver such as torch.linalg.solve is avoided because it costs O(D³) per timestep and requires materializing every H_t; CH is used instead, and is shown to be more numerically accurate and faster-converging than gradient descent, accelerated gradient descent, and conjugate gradient.

3. Hardware-aware chunk-wise kernels, including a chunk-wise backward pass. The authors derive a chunk-wise formula for ||H_t||_F (which requires back-propagating through a nested recurrence) and use implicit differentiation for the backward pass through the linear equations CH approximately solves. They state this is a first adoption of Chebyshev Iteration for training sequence-modeling layers in LLMs stably at scale.

4. A released implementation and model zoo. Code including Triton kernels for training and inference (vLLM) is released under Apache 2.0 at github.com/awslabs/hybrid-model-factory, along with a HuggingFace model zoo of GKA-based hybrid models at 8B and 32B scale.

Main Findings

• The KF objective is the distinguishing factor. The paper contrasts the Attention regression objective (which uses the full KV-cache, with weights that change at every timestep because the query changes) against SSM objectives like Gated DeltaNet's, which consider only the previous lossy state and the current token. The KF-based objective in Eq. 3 uses the entire past, with a weight η_i that is fixed once observed, and a regularizer λ·||S||²_F that controls the memorization capacity of the constant-sized state.

• Condition number bound from adaptive regularization. Setting λ_t = a·||H_t||_F yields κ_t ≤ (a+1)/a. The authors argue this matters because the worst-case numerical error of solving Eq. 3 is ε·κ, where ε ≈ 0.007 in bf16. They report that prior work [64] lower-bounds λ at 0.25 but has κ as large as 500, implying a worst-case error of 3.5, and that the regression objective in [66] has no regularization.

• CH converges more accurately than CG and is more stable. In a random regression setup with batch size 8, sequence length 2048, 8 heads, and head dimension 128, CG converges fastest within a few iterations, while CH reaches the same accuracy as CG at iteration 10 and eventually attains the smallest errors.

• Gradient stability favors CH over CG. CG with implicit differentiation as a standalone layer has gradients close to the exact solver up to a 10⁻³ relative difference, but in a 5-layer LLAMA where Attention is replaced, that difference is amplified to almost 1 — meaning CG completely deviates from the reference gradient. CH with implicit differentiation and CH with torch.autograd are eventually close to the exact solver up to a 10⁻⁶ difference, with the two curves nearly overlapping. Lemma 2 states that dq_t = dq̂_t for CH.

• The backward pass uses a full interaction mask. The intra-chunk backward term involves M_w = M diag(w_1, …, w_C) Mᵀ, which is generally a full matrix with no zero entries, unlike the triangular mask in the linear-SSM case. The authors state this allows all tokens to interact with all other tokens in the backward pass.

• MQAR results. Two-layer models per memory-layer type are trained on MQAR data and evaluated on a held-out test set, repeated across 4 learning rates spanning 10⁻⁴ to 10⁻². The paper presents accuracy-versus-model-dimension per sequence length and per number of key-value pairs in Fig. 3(a); the specific accuracy values are not included in the content available here, as the text is truncated at that point.

• Runtime comparison. Fig. 3(b) reports runtime of a single memory layer for one forward plus backward pass, with 8 heads, head dimension 128, batch size 4, averaged over 20 runs. The numerical runtimes are not reported in the content available here.

• Long-context gains. GKA achieves more than 10% relative improvement on long-context RAG and LongQA up to 128k tokens relative to SSM baselines, and outperforms existing SSM layers such as Mamba2 and Gated DeltaNet on short-context tasks. Benchmark suites named are RULER and HELMET for long context, LM-Harness for short context, and MQAR for synthetic recall.

• Vision transfer. GKAVision outperforms MambaVision on ImageNet classification, which the authors present as evidence that the recall advantage transfers beyond language.

• The β_t-augmented variant. A subsequent work [4] augments the forward recurrence with a learned input-selectivity gate β_t ∈ [0,1] modulating the per-token write into H_t and U_t. All experiments in this paper use β_t ≡ 1, but the released GKA implementation adopts the β_t-augmented form by default, and the authors recommend it as the canonical GKA layer going forward.

Methodology in Plain English

The starting point is a shift in what the layer is trying to do. Instead of updating a compressed state only from the current token and the previous state, the layer solves a regularized least-squares problem over every key-value pair seen so far. That problem — a weighted ridge regression — has a known closed form whose online solution is the Kalman Filter recursion. Because the KF recursion is inherently sequential, the authors rewrite the solution so it can be computed by solving (H_t + λ_t I)x = q_t and then multiplying by U_t, where H_t and U_t are running weighted covariances with a decay factor.

They cannot call a general linear solver directly: it costs O(D³) per timestep and forces every intermediate H_t to be materialized, which is expensive in memory traffic. So they approximate the solve with Chebyshev Iteration, an accelerated first-order method that needs only matrix-vector products, costs O(D²) per iteration, and converges in a few iterations. To keep the problem well-conditioned in bfloat16, they tie the regularization to the Frobenius norm of the data covariance, which caps the condition number. To keep the layer expressive, they replace fixed weights with learnable exponentially decaying gates γ_j ∈ [0,1], encoding a recency bias and giving linear-time implementation.

The engineering work is a chunked implementation: the sequence is split into chunks of size C, with N = T/C chunks, and only the initial state of each chunk is precomputed sequentially, so all within-chunk work runs in parallel without materializing every H_t. The Frobenius norm needed for the adaptive regularization is computed with a closed-form chunk-wise expression involving the Gram matrix of the keys. Because storing activations for backpropagation through the iterative solver is costly, the backward pass uses implicit differentiation and reuses CH on the upstream gradient; the authors prove this gives the same gradient as differentiating through CH directly.

Why This Matters

Impact on research. The paper gives a unifying lens on a large family of modern linear SSM layers — DeltaNet, Gated DeltaNet, Kimi Delta Attention — by identifying them as special cases of the Kalman Filter under an identity error-covariance assumption. That reframing suggests a principled axis for designing new layers (how much second-order information to retain) rather than tuning heuristics. It also makes a concrete numerical-stability argument for why regularized, non-exact solvers are needed when training such layers in bfloat16, and offers Chebyshev Iteration as a drop-in alternative to conjugate gradient with a proven gradient-equivalence property.

Real-world applications:

• Retrieval-augmented generation at long context. The paper reports over 10% relative improvement on long-context RAG up to 128k tokens, which is where the storage and compute cost of Attention is most painful.
• Long-document question answering. LongQA at 128k tokens is evaluated directly.
• Long-context assistants and agents. Any workload that must recall specific earlier spans (rather than summarize them) benefits from a memory that does not fade.
• Vision and multimodal classification. GKAVision outperforming MambaVision on ImageNet indicates the layer is not limited to language.
• Deployment on constrained hardware. Constant memory and linear compute in sequence length make the layer attractive where a growing KV-cache is infeasible.

Industry relevance. The work comes from AWS Agentic AI with an academic co-author, ships Apache-2.0 code, Triton kernels, vLLM inference support, and pretrained GKA-based hybrid models at 8B and 32B scale on HuggingFace — a production-oriented release rather than a paper-only artifact.

Future Directions

• Adopting the β_t-augmented recurrence. The paper's own experiments use β_t ≡ 1, yet the released implementation defaults to the β_t-augmented form, which the authors say consistently improves long-context performance with gains widening at longer contexts. Closing that gap between the evaluated and released variants is the most immediate next step.
• Scaling and throughput characterization. Section 5.2 on training throughput exists, but the throughput numbers are not present in the content available here; how GKA's chunk-wise kernels compare in wall-clock training cost to Mamba2 and Gated DeltaNet at larger scales is an open question for readers of the full paper.
• Extending the error-covariance treatment. The paper's central argument is that maintaining the full error covariance is what separates GKA from prior SSMs. How much of that benefit survives when the covariance is approximated or sparsified to reduce cost is not answered here.
• Beyond ImageNet for vision. The vision result is limited to ImageNet classification against MambaVision; whether recall advantages hold on detection, segmentation, video, or long-form multimodal tasks is untested in this work.

Target Audience

Researchers and engineers working on efficient Transformer alternatives, long-context modeling, and linear attention/SSM architectures. It is most useful to readers who are comfortable with state-space recurrences, regularized least squares, and numerical linear algebra, and who are deciding which memory layer to adopt or how to design a new one. Practitioners interested only in deployment may find the released code and model zoo more directly useful than the derivations.

Authors’ abstract

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.