Kronecker Delta Rule for Long Context Linear Attention

Anamaria-Roberta Hartl; Levente Zólyomi; Michael List; Günter Klambauer; Sebastian Böck; Sepp Hochreiter

cs.LGOct 19, 2026

Kronecker Delta Rule for Long Context Linear Attention

Authors: Anamaria-Roberta Hartl, Levente Zólyomi, Michael List, Günter Klambauer, Sebastian Böck, Sepp Hochreiter

Organizations: Institute for Machine Learning, Johannes Kepler University, Linz · NXAI GmbH, Linz

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.

Figures & tables

Figure 1: Representation of the memory matrix state revision for mLSTM, Gated DeltaNet, and KronLSTM. Interference with stored values: mLSTM affects all values, Gated DeltaNet affects those with overlapping keys, and KronLSTM affects only those with overlapping key and value pairs.

Figure 2: Stress test with high learning rates for the same mLSTM and KronLSTM architectures on language modeling. Learning becomes unstable for mLSTM as the gradient norm spikes show, while KronLSTM yields stable learning.

Figure 3: Architecture of the KronLSTM block.

Model	Hella	PIQA	ARC-e	ARC-c	Wino	LMB acc	Avg-6 $\uparrow$	LMB ppl $\downarrow$
Linear models
GDN	57.00	73.61	70.54	41.04	57.22	46.94	57.73	12.58
Mamba-3	57.30	73.01	70.29	40.19	57.46	45.64	57.31	12.95
KDA	57.48	73.39	71.84	41.55	57.14	47.35	58.12	12.67
GDN-2	57.45	72.96	71.13	41.89	57.30	46.92	57.94	12.54
mLSTM	56.05	73.07	71.30	41.30	56.99	44.38	57.18	13.87

Table 1: Downstream performance of 1.5B/100B models. Avg-6 is the mean accuracy over the five commonsense reasoning tasks and LAMBADA. Best per column in bold , second best underlined , separately for linear and hybrid models.

Model	SWDE	SQD	FDA	TQA	NQ	DROP	Avg-6 $\uparrow$
Linear models
GDN	28.62	37.33	14.79	60.49	16.47	21.18	29.81
Mamba-3	27.54	35.86	11.43	59.18	15.39	21.03	28.41
KDA	25.83	37.30	12.98	62.14	16.31	20.46	29.17
GDN-2	29.70	36.29	14.43	60.07	15.90	19.93	29.39
mLSTM	27.54	36.96	11.98	59.77	15.11	21.37	28.79

Table 2: Results for 1.5B models on the real-world recall suite, given in accuracies for 2k context length. Best per column in bold , second best underlined .

		Linear models						Hybrid models
Model	Len.	S1	S2	S3	MK	MQ	MV	S1	S2	S3	MK	MQ	MV
GDN	4k	100.0	96.0	54.8	28.0	44.8	46.8	51.8	52.6	21.8	31.0	40.6	38.3
	8k	100.0	45.6	13.0	24.0	21.2	23.5	26.6	31.2	22.2	24.6	17.8	18.8
	16k	99.2	17.8	5.4	17.8	8.6	13.8	12.6	16.8	3.6	18.4	6.9	8.8
Mamba-3	4k	92.0	68.2	26.4	26.8	20.0	18.4	51.6	56.0	30.2	34.4	19.1	21.3
	8k	59.0	27.0	5.4	20.6	4.2	3.6	26.6	31.2	15.0	24.2	10.7	22.4

Table 3: Results of 1.5B models on the NIAH tasks for 8k pretraining context and 4k, 8k, and 16k evaluation context. S1–S3: S-NIAH-1/2/3; MK: MK-NIAH-1; MQ: MQ-NIAH; MV: MV-NIAH. Best per task and length in bold , second best underlined , separately for linear and hybrid models.

Model	Avg-6 $\uparrow$	LMB ppl $\downarrow$	SWDE	SQD	FDA	TQA	NQ	DROP	Recall avg $\uparrow$
Linear models
GDN-2	58.75	8.36	29.61	36.39	19.60	60.19	14.54	19.36	29.95
KronLSTM	60.19	7.66	31.68	39.88	18.15	60.37	15.36	21.03	31.08
Hybrid models
GDN-2	59.32	8.27	40.41	40.95	30.40	62.50	15.24	21.75	35.21
KronLSTM	60.10	7.64	46.08	44.10	23.96	63.39	16.95	22.90	36.23

Table 4: Long-context (32k) extension results for GDN-2 and KronLSTM models. Avg-6 is the mean accuracy of the five reasoning tasks plus LAMBADA acc; the remaining columns give accuracy on the real-world recall suite evaluated at 2048 tokens. The best per column is in bold .

Figure 4: GDN-2 and KronLSTM results on NIAH for 32k context extension for linear and SWA-hybrid models.

Figure 5: State tracking results for Gated DeltaNet (GDN), Mamba-3, and KronLSTM.

Figure 6: Training throughput of KronLSTM against baselines on a single H100-80GB.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Definition	Symbol	Type
Scalars and dimensions
time step (token index)	$t$	$\mathbb{N}$
key and query dimension per head	$d_{k}$	$\mathbb{N}$
value dimension per head	$d_{v}$	$\mathbb{N}$
layer input dimension	$d$	$\mathbb{N}$
number of stored associations (App.)	$n$	$\mathbb{N}$

Appendix

Table A1: Notation used throughout the paper and the appendix.

Model	Hella	PIQA	ARC-e	ARC-c	Wino	LMB acc	Avg-6 $\uparrow$	LMB ppl $\downarrow$
GDN	41.28	66.05	57.87	31.48	49.64	32.31	46.44	34.36
Mamba-3	41.89	67.25	60.82	31.23	52.33	35.11	48.10	29.51
KDA	42.08	66.16	60.48	31.40	51.30	36.70	48.02	27.84
GDN-2	42.33	66.87	60.23	29.69	52.41	33.84	47.56	29.14
mLSTM	40.29	67.19	60.77	29.78	50.75	29.87	46.29	43.40
KronLSTM	43.09	67.03	61.03	32.00	52.41	35.05	48.43	28.28

Appendix

Table A2: Downstream performance of 400M/40B linear models. Avg-6 is the mean accuracy over the five commonsense reasoning tasks and LAMBADA. Best per column in bold , second best underlined .

Ablations	Avg-6 $\uparrow$	LMB ppl $\downarrow$	Recall $\uparrow$	NIAH-Avg@4k	@8k	@16k
base	59.72	11.716	30.3	59.5	30.5	14.3
no $\ell_{2}$ -norm on $\bm{v}_{t}$	59.41	11.844	31.3	39.6	32.7	12.6
$\beta$ on both sides	60.00	11.479	30.8	46.9	29.4	11.8
non-learnable target	58.17	12.495	30.1	41.9	26.2	10.1

Appendix

Table A3: Write-rule ablation at 1.5B/100B, relative to KronLSTM . Avg-6 is the mean of the five commonsense reasoning tasks plus LAMBADA accuracy. Recall is the six-task real-world recall suite at 2k context. RULER is averaged at each evaluation length. Best per column in bold , second best underlined .

Model	Len.	S1	S2	S3	MK	MQ	MV
Linear models
GDN-2	4k	100.0	74.0	68.2	28.4	56.5	62.5
	8k	100.0	55.6	17.6	23.2	38.8	41.9
	16k	100.0	10.2	7.4	9.4	7.4	8.9
	32k	97.8	2.6	0.8	3.0	1.5	7.2
KronLSTM	4k	100.0	99.6	92.6	30.8	37.1	43.9

Appendix

Table A4: RULER NIAH for the 32k-context-extended linear models, evaluated on identical documents. Better of the two per task and length in bold .

	Gram in the solve	right-hand side	coupling to $\bm{S}_{0}$	write of token $r$
DeltaNet	$\operatorname{diag}(\bm{\beta})\operatorname{stril}(\bm{K}\bm{K}^{\top})$	$\operatorname{diag}(\bm{\beta})\bm{V}$	$\bm{S}_{0}\bm{w}_{r}$ , $\;\bm{W}=\bm{\Lambda}\operatorname{diag}(\bm{\beta})\bm{K}$	$(\bm{u}_{r}-\bm{S}_{0}\bm{w}_{r})\,\bm{k}_{r}^{\top}$
Gated DeltaNet	same, $\times\,\gamma_{r}/\gamma_{s}$ (scalar)	same	same, decayed	same, decayed
KDA	$\operatorname{diag}(\bm{\beta})\,\bm{G}^{kk}$	$\operatorname{diag}(\bm{\beta})\bm{V}$	$\bm{S}_{0}\bm{w}_{r}$ , $\;\bm{W}=\bm{\Lambda}\operatorname{diag}(\bm{\beta})\widehat{\bm{K}}$	$(\bm{u}_{r}-\bm{S}_{0}\bm{w}_{r})\,\breve{\bm{k}}_{r}^{\top}$
Gated DeltaNet-2	$\operatorname{stril}\big((\bm{\gamma}\odot\bm{B}\odot\bm{K})\bar{\bm{K}}^{\top}\big)$	$\bm{W}\odot\bm{V}$	$\bm{S}_{0}\bm{y}_{r}$ , $\;\bm{Y}=\bm{\Lambda}(\bm{\gamma}\odot\bm{B}\odot\bm{K})$	$(\bm{u}_{r}-\bm{S}_{0}\bm{y}_{r})\,\breve{\bm{k}}_{r}^{\top}$
KronLSTM	$\operatorname{diag}(\bm{\beta})\,(\bm{G}^{kk}\odot\bm{G}^{vv})$	$\bm{c}$	$\big(\bm{\Lambda}\operatorname{diag}(\bm{\beta})\bm{p}\big)_{r}$ , $\;p_{r}=\bm{v}_{r}^{\top}\bm{S}_{0}\widehat{\bm{k}}_{r}$	$u_{r}\,\bm{v}_{r}\breve{\bm{k}}_{r}^{\top}$

Appendix

Table A5: Chunkwise forms of delta-rule memories and of KronLSTM, in the orientation

\bm{S}\in\mathbb{R}^{d_{v}\times d_{k}}

of this paper.

\bm{\Lambda}

is the inverse of the unit lower triangular matrix

\bm{I}+\bm{N}

built from the listed Gram; “coupling” is the term through which the incoming memory

\bm{S}_{0}

enters the write of token

r

;

\operatorname{stril}

is the strictly lower triangular part. For KDA and Gated DeltaNet-2 the decay-weighted key Gram is

\bm{G}^{kk}

of ( A152 ); for Gated DeltaNet-2,

\bm{b}_{r}\in[0,1]^{d_{k}}

and

\bm{w}_{r}\in[0,1]^{d_{v}}

are its channel-wise erase and write gates.

Delta-Rule Linear Attention

Delta-rule linear attention updates a compact recurrent memory to retrieve and revise sequence information with linear-time processing.

See all 35 papers

Kalman Delta Networks: Uncertainty-aware Associative Memory
Sep 7, 2026Ngoc Bui, Tinglin Huang, Rex YingMemory-Augmented Neural Networks Linear Attention
Department of Computer Science, Yale University
Dynamic Linear Attention
Jun 9, 2026Xin Wang, Hui Shen, Boyuan Zheng +7Long-Context Language Modeling Linear Attention
The Ohio State University · University of Michigan · ByteDance Seed
Triadic Linear Attention: Three-Dimensional Recurrent States for Long-Context Sequence Modeling
Sep 29, 2026Oliver Sieberling, Bharat Runwal, David Jin +3Linear Attention Long-Range Dependency Modeling
Massachusetts Institute of Technology · MIT-IBM Computing Research Lab

Sep 7, 2026J/K move · Enter open · S save

Kalman Delta Networks: Uncertainty-aware Associative Memory

Ngoc Bui, Tinglin Huang, Rex Ying

Department of Computer Science, Yale University

Memory-Augmented Neural Networks Linear Attention Delta-Rule Linear Attention

Open paper AskPDF arXiv

Linear attention enables efficient long-context inference by compressing token history into a fixed-size recurrent memory. This compression makes each update a trade-off between incorporating new information and preserving useful associations. Models such as DeltaNet, Gated DeltaNet, and KDA predict write strength from the current token representation, without explicitly tracking uncertainty in the stored memory. Yet this uncertainty matters: a new observation should have greater influence when the existing association is uncertain and less when it is already well supported. We introduce Kalman Delta Networks (KDNs), a family of linear-attention models that explicitly track memory uncertainty to guide each update. By formulating associative memory as a linear-Gaussian state-space model, KDNs propagate both the memory estimate and its uncertainty, using the Kalman gain to balance accumulated evidence against the reliability of new observations. This formulation also recovers standard delta-rule updates by replacing tracked covariance with a token-predicted isotropic surrogate. To support hardware-efficient training and inference, we derive Diagonal KDN and Isotropic KDN, which retain one uncertainty value per key channel and per head, respectively. Their uncertainty updates admit associative scans with logarithmic parallel depth, requiring only

O(d_k)

and

O(1)

auxiliary state per head. Across controlled pretraining at 750M and 1.3B parameters, both variants consistently improve perplexity and mean downstream accuracy over the evaluated state-of-the-art linear-attention baselines.

Abstract

Delta-Rule Linear Attention

Explore similar work

Kalman Delta Networks: Uncertainty-aware Associative Memory

Dynamic Linear Attention

Triadic Linear Attention: Three-Dimensional Recurrent States for Long-Context Sequence Modeling