Gated Slot Attention-2: Two-Sided Associative Memory Correction in Linear Attention

Ruijie Li; Shengnan Ding; Weimin Zhang; Derick Tang; Zhanpeng Zeng; Qinsong Zeng; Ming Chen; Jiaxi Hu

cs.LGOct 2, 2026

Gated Slot Attention-2: Two-Sided Associative Memory Correction in Linear Attention

Authors: Ruijie Li, Shengnan Ding, Weimin Zhang, Derick Tang, Zhanpeng Zeng, Qinsong Zeng, Ming Chen, Jiaxi Hu,

Organizations: The Hong Kong University of Science and Technology (Guangzhou) · Tencent

Abstract

Linear attention models have emerged as efficient alternatives to standard attention, but effectively managing their fixed-size recurrent memory remains challenging. To improve memory, recent work has explored two distinct directions: delta-rule variants for precise correction of values associated with keys, and slot-based architectures such as Gated Slot Attention for modeling key and value memories in two stages. We observe that these directions are complementary--the delta rule provides effective memory correction, while the two-stage structure provides a natural way to operate on both sides of an association. Building on this insight, we introduce a new Gated Oja Rule for key-side correction and extend it with decoupled erase and write control to obtain Gated Oja Rule-2. We then introduce Gated Slot Attention-2 (GSA2), which combines Gated Oja Rule-2 for key-side correction with Gated Delta Rule-2 for value-side correction through shared latent slots. We further derive a hardware-efficient chunkwise algorithm for parallel training. Experiments demonstrate that GSA2 consistently improves over strong linear-attention baselines across benchmarks while retaining linear-time sequence modeling and constant-memory recurrent decoding.

Figures & tables

Figure 1: Architecture of GSA2 . Left: the hybrid backbone interleaves GSA2, MLP, and SWA blocks. Right: the GSA2 token mixer. The query and key branches apply linear projection, short convolution, SiLU, and L2 normalization, while the value branch follows the same pipeline without L2 normalization. The shared-slot branch uses a low-rank two-layer projection followed by short convolution, SiLU, and L2 normalization. Gated Oja Rule-2 and Gated Delta Rule-2 perform key- and value-side correction through shared latent slots, with separate decay, erase, and write gates controlling the two recurrent memories. The output is RMS-normalized, modulated by a sigmoid gate from a low-rank two-layer projection, and linearly projected.

Recurrent models
Model	Real-world Retrieval (2K)
Model	SWDE	SQD	FDA	TQA	NQ	DROP	Avg.
	acc $\uparrow$	acc $\uparrow$	acc $\uparrow$	acc $\uparrow$	acc $\uparrow$	acc $\uparrow$	acc $\uparrow$
Mamba-2 ∗	17.24	32.38	14.53	58.35	18.91	19.60	26.84
Gated DeltaNet ∗	17.90	32.67	18.52	59.60	20.16	19.69	28.09
KDA ∗	22.49	35.10	14.90	58.12	19.58	21.80	28.67

Table 3: Real-world retrieval at 2K context. All non-Transformer hybrid models use SWA.

Figure 2: RULER accuracy of recurrent and attention/hybrid models. Each task is averaged over 1K–8K contexts. NIAH Avg. summarizes the 24 S-NIAH and MK-NIAH settings, while Cross-Task Avg. summarizes the 28 settings from MQ, MV, CWE, FWE, HPQA, SQuAD, and VT. Colors denote base models; hatched and solid bars indicate recurrent and attention/hybrid configurations, respectively.

Table 4: LongBench performance across five task categories. Avg. is the mean over all 11 tasks. All non-Transformer models in the hybrid group use SWA.

Ablation configuration		Language modeling		Common.	RULER			Real-world retrieval	LongBench
Stage 1	Stage 2	Wiki. $\downarrow$	LMB. $\downarrow$	avg $\uparrow$	S-NIAH $\uparrow$	MK-NIAH $\uparrow$	Rem. $\uparrow$	avg $\uparrow$	avg $\uparrow$
Coupled	Coupled	16.22	10.29	53.25	68.2	9.9	30.0	29.60	17.1
Coupled	Decoupled	16.01	10.65	52.46	68.0	10.9	24.7	31.26	14.8
Decoupled	Coupled	16.28	11.25	53.37	65.3	9.8	28.3	30.17	16.2
Decoupled	Decoupled	15.69	10.68	54.08	72.8	11.4	31.8	31.51	17.3
Gated Delta Rule-2 in both updates		15.92	10.27	53.03	72.0	13.1	29.5	29.76	16.8

Table 5: Ablation study of GSA2. Common. averages LAMBADA accuracy and nine zero-shot common-sense reasoning tasks. S-NIAH, MK-NIAH, and Rem. are averages over the corresponding RULER task groups. Real-world retrieval averages six tasks at 2K context, and LongBench averages 11 tasks.

Panel of Figure 3 — Figure 3: Training throughput of parameter-matched 1.3B models on a single NVIDIA H20 GPU under a fixed 16K-token budget per optimization step. Sequence length increases from 2K to 16K while batch size decreases correspondingly from 8 to 1. The throughput evaluation protocol is detailed in Appendix A.3.3 .

Linear Attention

Linear attention reduces sequence-length costs by replacing pairwise token interactions with kernel features, recurrent states, or compact memory.

See all 129 papers

Gated DeltaNet-2: Decoupling Erase and Write in Linear Attention
May 21, 2026Ali Hatamizadeh, Yejin Choi, Jan KautzLong-Context Language Modeling Recurrent Transformers
NVIDIA
Kaczmarz Linear Attention
May 9, 2026Jiaxuan Zou, Ruifeng Ren, Yong LiuLong-Context Language Modeling Recurrent Transformers
School of Mathematics and Statistics Xi’an Jiaotong University Xi’an, China · Gaoling School of Artificial Intelligence Renmin University of China Beijing, China
Variational Linear Attention: Stable Associative Memory for Long-Context Transformers
May 11, 2026Vishal Pandey, Gopal SinghEfficient Attention Linear Attention
Independent Researcher London, UK · Metriqual Athens, GR

May 21, 2026J/K move · Enter open · S save

Gated DeltaNet-2: Decoupling Erase and Write in Linear Attention

Ali Hatamizadeh, Yejin Choi, Jan Kautz

NVIDIA

Long-Context Language Modeling Recurrent Transformers Linear Attention Delta-Rule Linear Attention Gated Delta Networks

Open paper AskPDF arXiv

Linear attention replaces the unbounded cache of softmax attention with a fixed-size recurrent state, reducing sequence mixing to linear time and decoding to constant memory. The hard part is not just what to forget, but how to edit this compressed memory without scrambling existing associations. Delta-rule models subtract the current read before writing a new value, and Kimi Delta Attention (KDA) sharpens forgetting with channel-wise decay. But the active edit still uses a single scalar gate to control two different things: how much old content to erase on the key side and how much new content to commit on the value side. We introduce Gated DeltaNet-2, which generalizes both Gated DeltaNet and KDA by inheriting adaptive forgetting and channel-wise decay while addressing their shared limitation, the scalar tie between erasing and writing. Gated Delta Rule-2 separates these roles with a channel-wise erase gate b_t and a channel-wise write gate w_t, reducing to KDA when both gates collapse to the same scalar and to Gated DeltaNet when the decay also collapses. We derive a fast-weight update view, a chunkwise WY algorithm with channel-wise decay absorbed into asymmetric erase factors, and a gate-aware backward pass that preserves efficient parallel training. At 1.3B parameters trained on 100B FineWeb-Edu tokens, Gated DeltaNet-2 achieves the strongest overall results among Mamba-2, Gated DeltaNet, KDA, and Mamba-3 variants across language modeling, commonsense reasoning, and retrieval. Its advantage is most pronounced on long-context RULER needle-in-a-haystack benchmarks, where it improves the evaluated multi-key retrieval setting and remains strong in both recurrent and hybrid settings. Code is available at https://github.com/NVlabs/GatedDeltaNet-2.

Abstract

Linear Attention

Explore similar work

Gated DeltaNet-2: Decoupling Erase and Write in Linear Attention

Kaczmarz Linear Attention

Variational Linear Attention: Stable Associative Memory for Long-Context Transformers