🤖 machine learning

Kaczmarz Linear Attention

The paper introduces Kaczmarz Linear Attention (KLA), a modified Gated DeltaNet that replaces its empirically learned update coefficient with a theoretically derived, key-norm-normalized Kaczmarz step size, resulting in superior perplexity, long-context stability, and decoding efficiency without altering the model's architecture or state shape.

Original authors: Jiaxuan Zou, Ruifeng Ren, Yong Liu

Published 2026-05-12

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Jiaxuan Zou, Ruifeng Ren, Yong Liu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to read a very long book. The robot needs to remember what it read earlier to understand the current sentence.

The Problem: The "Quadratic" Bottleneck
Traditional AI models (Transformers) work like a student who, for every new word they read, has to flip back through the entire book to check every previous word to see how they connect. If the book is short, this is fine. But if the book is 100,000 pages long, the student has to do a massive amount of work for every single word. This gets so slow and expensive that it's practically impossible to scale up.

The Solution: The "Recurrent State"
Newer models try to fix this by acting like a student with a small, fixed-size notebook. Instead of flipping back to the whole book, they just update their notebook as they read. They write down the most important bits, forget the rest, and keep moving. This is fast (linear time), but it's hard to get right: What should they write? How much should they erase? And how should they update the note if they see the same topic again?

The Previous Attempt: Gated DeltaNet (GDN)
One popular model, called Gated DeltaNet (GDN), uses a "notebook" approach. When it sees a new piece of information, it calculates the difference between what it thinks it knows and what it actually sees, then writes that difference into the notebook.

However, GDN has a flaw: it uses a "learned guess" (a number it figures out during training) to decide how big of a change to make. It's like a student guessing, "Hmm, I think I should write this down with a size 5 marker." Sometimes they use a marker that is too big (smudging the page), and sometimes too small (the writing is faint and gets lost). This guess is just a habit the model learned, not a mathematical rule.

The New Idea: Kaczmarz Linear Attention (KLA)
The authors of this paper, Jiaxuan Zou and colleagues, asked: "Can we stop guessing and use math to decide exactly how big the change should be?"

They looked at an old math method called the Kaczmarz projection.

The Analogy: Imagine you are trying to draw a line on a piece of paper that passes through a specific point. You have a ruler (your current state). If your ruler doesn't hit the point, you need to nudge it.
The Insight: The Kaczmarz method says the best way to nudge the ruler is to measure how "strong" or "loud" the point is. If the point is very loud (a strong signal), you only need a tiny nudge to hit it. If the point is quiet (a weak signal), you need a big nudge.

In the paper's language, they look at the "Key" (the signal) and measure its strength (its "norm"). They then calculate a precise step size:

Step Size = (Learning Rate) / (Strength of the Signal)

This is the Kaczmarz Coefficient.

What Changed?
The authors didn't build a new robot or a new notebook. They didn't change the hardware. They simply swapped the "guessing" number in the GDN model for this precise, mathematically derived number.

Old Way: "I'll write this down with a marker size of 0.5 because my training data told me to."
New Way (KLA): "I'll write this down with a marker size of 0.5 divided by how loud this signal is."

The Results
Because this new rule is mathematically perfect for the specific job of updating the memory, the model performs better:

Smarter: It predicts the next word in a sentence more accurately (lower "perplexity") than the previous best models.
Longer Memory: It can handle much longer contexts (up to 65,000 words) without getting confused or forgetting things, whereas the old model started to struggle.
Better at Tasks: In tests where the model had to find a specific "needle" in a huge "haystack" of text, KLA got 100% correct, while others missed.
Just as Fast: Because they only changed the math formula for the update and not the structure of the notebook, the model runs just as fast as the old one. In fact, it decodes (generates text) 2.1 times faster at long lengths.

In Summary
The paper introduces KLA, a model that keeps the same fast, efficient structure as its predecessor but replaces a "guessing" update rule with a precise, mathematically derived one. It's like taking a car that already drives well and swapping out the driver's guesswork for a perfect GPS navigation system. The car is the same, but it gets to the destination more accurately and efficiently.

Technical Summary: Kaczmarz Linear Attention (KLA)

Problem Statement

Long-context language modeling is hindered by the quadratic computational and memory costs of standard Softmax attention in Transformers. While linear-time recurrent models address this by compressing context into a fixed-size state, the design of the state-update rule becomes the critical bottleneck. Existing approaches, such as Gated DeltaNet (GDN), utilize a "delta-rule" to write the residual between a target value and the current state prediction, thereby editing existing associations rather than just accumulating them.

However, GDN relies on a learned scalar coefficient to balance forgetting and updating. The authors argue that this coefficient is an empirical design choice rather than a theoretically derived step size. An inappropriate coefficient can lead to suboptimal update magnitudes: over-correcting for large-norm keys and under-correcting for small-norm keys, resulting in state corruption and suboptimal retrieval performance.

Methodology: Kaczmarz Linear Attention

The paper proposes Kaczmarz Linear Attention (KLA), a modification of GDN that replaces the learned scalar with a theoretically derived step size based on the Kaczmarz projection method.

Theoretical Derivation

The authors frame the key-value write operation as a constrained optimization problem. For a token $t$ , the goal is to find the state $S_t$ closest to the decayed previous state $\tilde{S}_{t-1}$ that satisfies the linear constraint $S_t^T k_t = v_t$ (where $k_t$ is the key and $v_t$ is the value).

Optimization Objective: $\min_S \frac{1}{2}\|S - \tilde{S}_{t-1}\|_F^2$ subject to $S^T k_t = v_t$ .
Solution: The optimal solution is the orthogonal projection of $\tilde{S}_{t-1}$ onto the constraint hyperplane. This yields a step size (coefficient) of $\beta_t = 1/\|k_t\|_2^2$ .

The KLA Update Rule

KLA generalizes this exact projection into a practical update rule that includes a learnable relaxation gate $\eta_t$ and a numerical stabilizer $\epsilon$ :
$\beta_t = \frac{\eta_t}{\|k_t\|_2^2 + \epsilon}$
The state update becomes:
$S_t = \tilde{S}_{t-1} + \beta_t k_t e_t^T$
where $e_t = v_t - \tilde{S}_{t-1}^T k_t$ is the residual error.

Key Properties

The paper establishes that this coefficient satisfies three desirable properties simultaneously (in the unrelaxed case where $\eta_t=1, \epsilon=0$ ):

Orthogonal Projection: It performs the minimum-norm change to the state to satisfy the constraint.
Exact Line Search: It is the exact minimizer of the per-token loss along the update direction.
Normalized Gradient Step: It acts as a normalized Stochastic Gradient Descent (SGD) step.

Implementation

KLA is designed as a drop-in replacement for GDN. It preserves:

The state shape ( $S \in \mathbb{R}^{d_k \times d_v}$ ).
The gating mechanisms (decay and update gates).
The linear recurrence structure.
The chunkwise parallel training algorithm.

The only structural change is the computation of the diagonal coefficient matrix $B = \text{Diag}(\beta_1, \dots, \beta_C)$ within the chunkwise solver. No new hardware kernels or state layouts are required.

Key Contributions

Self-Normalized Delta Rule: Derivation of the KLA coefficient from the Kaczmarz projection, proving it is an orthogonal projection, exact line-search minimizer, and normalized gradient step for the per-token loss.
Drop-in Recurrent Algorithm: Demonstration that KLA retains the GDN recurrence and chunkwise solver, requiring only a change in the diagonal coefficient matrix.
Empirical Gains: Validation that KLA achieves lower perplexity and better state-tracking than GDN and other linear-time baselines without architectural expansion.
Preserved Efficiency: Confirmation that KLA maintains $O(n)$ time/memory complexity and matches GDN's prefill latency while improving decode throughput.

Experimental Results

Experiments were conducted at the 0.4B parameter scale with a 1B-token budget on the SlimPajama-672B dataset.

Language Modeling

Validation Perplexity (PPL): KLA achieved 8.09 PPL, outperforming GDN (8.50), DeltaNet (8.26), GLA (8.42), Longhorn (8.79), and Mamba2 (9.41).
Long-Context Extrapolation: Without long-context fine-tuning, KLA maintained the lowest perplexity across context lengths from 1K to 65K tokens, showing a consistent 3.7–4.1 PPL advantage over GDN.

Synthetic Tasks

Single-Needle-in-a-Haystack (S-NIAH): KLA reached 100% accuracy across all tested context lengths (1K–8K), whereas GDN dropped slightly and Mamba2 failed significantly.
Multi-Query Associative Recall (MQAR): KLA improved 8× length extrapolation accuracy by 7.03 points over GDN (73.84% vs 66.81%).
Stack & Palindrome: KLA showed competitive or superior performance in state-tracking and reversal tasks.

Efficiency

Prefill Latency: KLA matched GDN's prefill latency (e.g., ~1220ms at 131K context), confirming the Kaczmarz coefficient calculation does not impact the dominant chunkwise kernel cost.
Decode Throughput: At 32K context, KLA achieved 3.54 tokens/s, which is 2.1× higher than GDN (1.65 tokens/s) and 1.7× higher than Mamba2.

Ablation Studies

The authors ablated the normalization strategy, gating mechanism, and state expansion. Results indicated that the key-norm normalization is critical for stability; removing it or using a learned scalar alone led to significant degradation in pretraining stability (PPL) and synthetic task performance.

Significance and Claims

The paper claims that the key-norm-normalized Kaczmarz coefficient is a "first-order design axis" for delta-rule sequence models. By treating the write operation as a constrained projection, KLA resolves the issue of suboptimal update magnitudes inherent in empirically learned scalars.

The significance lies in the fact that these improvements are achieved without changing the recurrent state, hardware kernel, or architectural complexity. The authors posit that the coefficient itself, rather than the state capacity or gating structure, is the primary factor determining whether a delta-rule write successfully solves its local regression problem. This suggests that future linear-time models can achieve substantial gains in accuracy, extrapolation, and decoding efficiency simply by adopting this theoretically grounded step size.