Kronecker Delta Rule for Long Context Linear Attention

AI summary

Overview

Research area: Machine learning — sequence modeling architectures, specifically linear attention and recurrent associative-memory models as alternatives to Transformer softmax attention.

Technical level: Advanced. The paper is written for readers comfortable with matrix algebra, gradient-based learning theory, and LLM architecture internals; the core intuition, however, is explainable without the proofs.

Scope: The paper introduces a new memory-update rule (the Kronecker Delta Rule) for associative matrix memories, embeds it in a gated linear attention architecture called KronLSTM, proves bounds on its memory norm, and benchmarks 1.5B-parameter models against Gated DeltaNet, Mamba-3, Kimi Delta Attention, Gated DeltaNet-2, and mLSTM on language modeling, retrieval, state tracking, and throughput.

What This Paper Is About

Softmax attention in Transformers gets quadratically more expensive as context grows, so researchers replace it with linear attention that squeezes the context into a fixed-size memory matrix updated token by token. The problem is that existing update rules either let the memory grow unboundedly (causing numerical instability) or, when they correct stored content, disturb associations the model has already learned. The paper proposes an update rule that fixes only one scalar association at a time along a single direction, which reduces this interference, and builds an architecture around it.

Key Contributions

1. The Kronecker Delta Rule. A new update rule for associative matrix memories that corrects the scalar association between a key and a value toward an input-dependent target, changing the memory only along the outer product of that specific key–value pair rather than along a memory-dependent direction.

2. KronLSTM. A linear attention architecture embedding the rule with three separate controls: a vector forget gate acting on the key dimension, a sigmoid erase gate, and a sigmoid write gate (plus an output gate and per-head RMSNorm). It carries no normalizer or stabilizer state, unlike mLSTM.

3. Stability proofs. Mathematical results showing the expected memory converges to a unique least-squares steady state and that the expected memory norm is bounded at every step and in the stationary distribution, guaranteeing stable learning and gradient flow.

4. Empirical validation at 1.5B parameters. Experiments on 100B tokens of FineWeb-Edu covering commonsense reasoning, real-world recall, RULER needle-in-a-haystack tasks up to 16k (and 32k after context extension), state tracking, and training throughput, each run both as a pure linear model and as a hybrid with sliding-window attention.

Main Findings

• Best average language-modeling accuracy. On the six-task average (five commonsense reasoning tasks plus LAMBADA), linear KronLSTM reaches 59.72 against the strongest competitor KDA at 58.12 — a 1.6 accuracy point gain — and the lowest LAMBADA perplexity at 11.72 versus 12.54 for GDN-2. In the hybrid setting KronLSTM scores 59.56 with perplexity 11.18.

• Strongest real-world recall average. On the 2k-context recall suite (SWDE, SQuAD, FDA, TriviaQA, NQ, DROP), KronLSTM has the best average in both settings: 30.31 for linear and 37.40 for hybrid, the latter improving on the next-best mLSTM at 36.37.

• Advantage grows with context length. On RULER NIAH, hybrid KronLSTM has the highest 16k accuracy on every task. On multi-query NIAH, its lead grows from 2.4 to 4.1 accuracy points when going from 4k to 16k; on multi-value NIAH it trails GDN by 3.5 points at 4k but leads by 1.1 and 3.8 points at 8k and 16k. Hybrid KronLSTM outranks GDN-2 and Mamba-3 by 15.2 accuracy points on S2 at 4k.

• One clear failure case: task S1. Linear KronLSTM scores only 66.2 on S1 at 4k versus 100.0 for GDN and GDN-2, which the authors attribute to residual accumulation for recurring keys in a task they describe as the most artificial, with frequently repeated noise sentences.

• Hybrid models retrieve mainly inside the window — except KronLSTM. Baseline hybrid linear models reach at most 51.8, 26.6, and 12.6 accuracy on S1 at 4k, 8k, and 16k, closely matching the fraction of context covered by the attention window; KronLSTM exceeds these and is the only one retrieving needles beyond the sliding window.

• Lead survives 32k context extension. After branching at 75B tokens and training the remaining 25B at 32k sequence length, linear KronLSTM scores 60.19 Avg-6 and 7.66 LAMBADA perplexity versus GDN-2 at 58.75 and 8.36; it wins five of six recall tasks. Perplexity improves by about four points for both models from extension.

• Competitive state tracking. Using an erase gate scaled to (0, 2), KronLSTM matches Gated DeltaNet on parity and S₃ permutation (both ahead of Mamba-3), while Mamba-3 wins on modulo arithmetic.

• Faster than GDN-2 but slower than the cheapest baselines. On a single H100 (80 GB) with 16,384 tokens per step, KronLSTM reaches 36.8–44.5K tokens/s, 11–18% more than GDN-2 at every shape and within 4–8% of KDA up to 8K×2; GDN, with scalar decay, remains fastest.

• Bounded memory, stable at high learning rates. A stress test with high learning rates shows mLSTM's gradient norm spiking while KronLSTM trains stably. With ℓ₂-normalized keys and values and a sigmoid erase gate, every update is non-expansive, and the expected squared memory norm is bounded by four times the expected target energy divided by the smallest eigenvalue of the key–value second-moment matrix, with the transient factor decaying geometrically from 4 to 1.

Methodology in Plain English

The starting point is a simple picture of linear attention: the model keeps a matrix that acts like a lookup table, and each new token writes a key–value association into it. The classic rule just adds new associations on top of old ones, so the table grows and old entries get buried. The Delta Rule instead subtracts what is currently stored before writing the new value — but because the correction direction depends on the table's current contents, it also nudges other entries that share similar keys.

The authors' alternative measures how strongly the memory already links the current key to the current value, producing a single number, and corrects only that number toward a target. The correction direction is the outer product of the key and value alone, so it never depends on the rest of the memory. They show this is equivalent to the classic Widrow–Hoff Delta Rule applied to the vectorized outer product of key and value, which also means an existing efficient chunkwise-parallel training scheme carries over. They then analyze a long stream of random key–value–target triples and prove the expected memory converges to the least-squares solution and stays bounded.

For the architecture, the step size and target become sigmoid gates computed from the input, and the memory is decayed by a diagonal forget gate operating on the key dimension, letting each key direction decay separately. A learned per-head scalar exponent rescales each head's retention horizon, so different heads can hold information over very different timescales. Keys, queries, and values are ℓ₂-normalized, and the read-out passes through a sigmoid output gate and per-head RMSNorm.

Experiments train 1.5B-parameter models on 100B tokens of FineWeb-Edu with the Llama2 tokenizer at 8k context, comparing against five recurrent linear attention baselines in both pure-linear and 5:1 sliding-window-hybrid configurations. A separate experiment extends the training context to 32k on a long-document mixture derived from ProLong-64K.

Why This Matters

Impact on research. The paper reframes memory interference in linear attention as a design variable rather than a fixed property of the Delta Rule family, and provides a capacity argument: because each triple imposes one linear constraint on the d_k·d_v entries of the memory, up to d_k·d_v scalar associations can be stored exactly, against only d_k value vectors for the standard Delta Rule — a single key can hold several orthogonal values. It also adds a data-distribution-dependent bound on memory norm that does not rely on ratio-based normalization.

Real-world applications:

• Long-document question answering and retrieval over hundreds of pages, where the paper's recall and NIAH results are directly relevant.
• Repository-scale code assistance, where context is long and memory footprint must stay constant.
• Conversational agents and personal assistants that must retain facts over very long interactions without a growing key–value cache.
• Streaming or memory-constrained inference on long sensor, log, or genomic sequences, where a fixed-size state is a hard requirement.

Industry relevance. Constant memory footprint with respect to context length, training throughput within a few percent of the fastest linear attention kernels, and drop-in compatibility with a Transformer block (analogous to mLSTM and sLSTM in xLSTM) make KronLSTM a practical candidate for frontier-model hybrid stacks. The measured 11–18% throughput advantage over GDN-2 at matched shapes is a concrete training-cost argument.

Future Directions

• Scaling KronLSTM up beyond the 1.5B-parameter, 100B-token regime reported here, and evaluating the 400M-scale results referenced in the appendix as a complementary data point.
• Hybrids containing full softmax attention layers, which the authors list explicitly as future work, in contrast to the sliding-window hybrids tested here.
• Fixing the S1 weakness. The residual accumulation for recurring keys that the authors identify as the cause of KronLSTM's S1 deficit is an open architectural question.
• Extending beyond 32k. Both KronLSTM and GDN-2 drop substantially in retrieval accuracy at 32k relative to 16k, so how the interference advantage behaves at much longer contexts remains untested.

Target Audience

Researchers and engineers working on efficient sequence modeling, linear attention, and recurrent alternatives to Transformers, particularly those designing hybrid architectures for long-context LLMs or implementing fast training and inference kernels. The theoretical sections on convergence and memory-norm bounds will also interest readers studying the stability of gated associative memories. Readers seeking only an architectural intuition will find the rule definition, the interference property, and the benchmark tables sufficient.

Authors’ abstract

We propose KronLSTM, a linear attention architecture based on the Kronecker Delta Rule, a novel update rule for associative matrix memories. KronLSTM extends the xLSTM family with a memory-revision mechanism that corrects targeted associations while modifying the memory only along a single direction in matrix space. This minimally invasive memory revision reduces interference with previously stored information compared with delta-rule updates, making it particularly well suited for long-context language modeling. The resulting memory dynamics can be mathematically described as a Markov chain that converges to a steady-state distribution, while the expected norm of the memory matrix remains uniformly bounded throughout the entire input sequence. Hence, KronLSTM ensures stable memory dynamics without requiring additional normalization and stabilization mechanisms for the memory state. Experiments show that KronLSTM is competitive with, and often outperforms, existing linear attention architectures, including models of the (Gated) DeltaNet family, with particularly strong performance on long context benchmarks. These results establish the Kronecker Delta Rule as an effective and analytically tractable foundation for stable associative memory in recurrent linear attention architectures. Our code (https://github.com/NX-AI/KronLSTM) and models (https://huggingface.co/collections/NX-AI/kronlstm) are open-source.