PaperlayerView exact revision on arXiv

Source-linked AI summary

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, Yuxuan Liang

arXiv:2610.02816v1cs.LG

TL;DR

Linear attention offers efficient fixed-size recurrent memory, but organizing and updating that memory remains challenging. GSA2 combines shared-slot key- and value-side correction with chunkwise training, and reports strong performance across benchmarks, particularly on RULER and LongBench.

  • Problem

    Linear attention compresses sequence history into fixed-size recurrent memory, but effectively organizing and updating that memory remains challenging.

    Source: §1 · p. 1

  • Method

    GSA2 combines Gated Oja Rule-2 key-side correction and Gated Delta Rule-2 value-side correction through shared latent slots.

    Source: §1 · p. 2

  • Results

    At 1.3B scale, GSA2 achieves strong performance across language modeling, commonsense reasoning, synthetic and real-world retrieval, and general long-context evaluation, with particularly consistent gains on RULER and LongBench.

    Source: §6 · p. 10

  • Takeaways & Limitations

    GSA2 retains fixed-size recurrent memory and efficient chunkwise training while providing complementary associative correction on both sides of shared latent-slot memory.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

Linear attention compresses sequence history into fixed-size recurrent memory, but organizing and updating that memory remains challenging. GSA2 combines two-stage slot memory with complementary key- and value-side correction.

1 INTRODUCTION

  • GSA retains separate key and value memories over shared latent slots, while delta-rule models provide explicit correction that primarily acts on values.GSA (Zhang et al., 2024) and delta-rule approaches including DeltaNet (Yang et al., 2024) illustrate the complementary directions.Source: §1 · pp. 1–2
  • GSA2 combines Gated Oja Rule-2 for key-side correction with Gated Delta Rule-2 for value-side correction through shared latent slots.The first stage associates keys with slots; the second associates those same slots with values.Source: §1 · p. 2
  • Gated Oja Rule adapts Oja’s normalized Hebbian update to provide key-side correction, and decoupling erase and write operations yields Gated Oja Rule-2.GSA2 combines it with Gated Delta Rule-2, whose erase and write controls are also decoupled.Source: §1 · pp. 1–2
  • The chunkwise formulation supports efficient parallel training while retaining linear-time sequence modeling and constant-size recurrent memory.The method reuses Gated Delta Rule-2’s chunkwise formulation and derives a corresponding WY/UT-style formulation for Gated Oja Rule-2.Source: §1 · p. 2
  • GSA2 performs strongly across language modeling, commonsense reasoning, retrieval, and long-context benchmarks, with particularly clear gains on RULER and LongBench.It preserves favorable recurrent length scaling with only a modest constant throughput overhead.Source: §1 · p. 2

2 PRELIMINARIES

Delta-rule memory corrects values associated with keys, with later variants adding adaptive forgetting and decoupled erase–write control. Gated Slot Attention instead factors associations through shared latent slots, but uses gated accumulation rather than explicit error correction.

2.1 DELTA-RULE ASSOCIATIVE MEMORY

  • DeltaNet corrects the value associated with the current key using the value-prediction error under that key.The update strength is controlled by β_t; the correction is an online gradient step.Source: §2.1 · p. 2
  • GDN2 decouples erasing old content from writing new content, using separate gates for the two operations.Under this paper’s state orientation, b_t controls erasing along the key-side axis and c_t controls writing along the value-side axis.Source: §2.1 · p. 3

2.2 GATED SLOT ATTENTION AS TWO-STAGE ASSOCIATIVE MEMORY

  • GSA factorizes a key–value association through M latent slots, maintaining slot–value associations in one of its memories.A query is first mapped to latent-slot space and then decoded through the value memory.Source: §2.2 · p. 3
  • GSA provides an intermediate slot space and compact recurrent memory, but updates both memories through gated accumulation rather than explicit error correction.This contrast motivates complementary key- and value-side corrections.Source: §2.2 · p. 3

3 GATED SLOT ATTENTION-2

GSA2 combines key-side Gated Oja Rule-2 and value-side Gated Delta Rule-2 within GSA’s two-stage slot memory. The section develops the Oja updates, shared-slot mechanism, chunkwise training, and model designs.

3 GATED SLOT ATTENTION-2

  • GSA2 combines Gated Oja Rule-2 and Gated Delta Rule-2 through shared latent slots to correct both sides of an association.The shared slot couples key–slot and slot–value memories.Source: §3 · p. 3, §3.2 · p. 5

3.1 GATED OJA RULE: KEY-SIDE CORRECTION

  • OJA2 provides key-side correction and separates erase and write control, with the write gate selecting key channels.When its erase and write gates are set to the shared scalar βt, the update reduces to the Gated Oja Rule.Source: §3.1 · pp. 3–4

3.2 TWO-STAGE MEMORY WITH SHARED SLOTS

  • The two shared-slot stages apply Oja correction to key–slot associations and Delta correction to slot–value associations, with independent control gates.GSA2 uses SiLU on corrected intermediate slots rather than GSA’s softmax, so the slots can contribute independently.Source: §3.2 · pp. 4–5

3.3 HARDWARE-EFFICIENT CHUNKWISE TRAINING

  • A WY/UT-style chunkwise formulation makes intra-chunk computation dominated by dense matrix multiplications.The method reuses GDN2’s chunkwise formulation for the Delta stage and derives the corresponding formulation for Oja updates.Source: §3.3 · pp. 5–6

3.4 BLOCK DESIGN AND MODEL FAMILIES

  • The recurrent model alternates GSA2 mixers with MLP layers, while the hybrid model interleaves GSA2, MLP, and SWA blocks.With a fixed SWA window, the hybrid model has linear sequence scaling and a bounded attention cache.Source: §3.4 · p. 6

4 EXPERIMENTS

Across language modeling, reasoning, retrieval, and long-context benchmarks, GSA2 performs strongly, with especially clear gains on RULER and LongBench. Ablations favor complementary key- and value-side correction, while throughput results show favorable sequence-length scaling with a moderate constant cost.

Language modeling and common-sense reasoning

  • GSA2 achieves the best common-sense average in both recurrent and hybrid settings while remaining strong on language modeling.Its gains persist against strong recurrent baselines and when combined with SWA.Source: §4 · p. 7

In-context retrieval on real-world tasks

  • GSA2 achieves the best average on real-world retrieval in the recurrent setting; OJA and OJA2 are particularly competitive in the hybrid setting.These tasks use 2K inputs, matching the SWA window, so local attention can directly access most relevant evidence.Source: §4 · p. 7, Table 3 · p. 7

In-context retrieval on synthetic data

  • GSA2 achieves the strongest overall RULER retrieval performance in both recurrent and hybrid settings, outperforming GSA, GDN2, and OJA2.The reported gains support combining key-side Gated Oja Rule-2 with value-side Gated Delta Rule-2.Source: §4 · p. 7, Figure 2 · p. 8

Long-context evaluation

  • GSA2 achieves the highest overall LongBench average in both recurrent and hybrid settings.LongBench covers five task categories and averages 11 tasks.Source: §4 · p. 7, Table 4 · p. 8

Ablation studies

  • Pairing key-side Oja with value-side Delta correction outperforms applying the same correction rule to both memories.Standalone Gated DeltaNet-2 generally outperforms OJA2, so Oja's benefit is its complementary key-side role within the two-stage memory.Source: §4 · pp. 7–8
  • Decoupling erase and write in both memories gives the strongest overall performance, while decoupling only one side yields less consistent gains.The default 128-slot configuration provides the best overall balance across evaluated tasks.Source: §4 · p. 8

Throughput comparison

  • At 2K to 16K sequence lengths, GSA2 throughput falls from 11.82 to 11.32 Kt/s (4.2%), versus the Transformer's 14.37 to 8.60 Kt/s (40.2%).The evaluation uses parameter-matched 1.3B models on one H20 GPU under a fixed 16K-token budget per step.Source: §4 · pp. 8–9, Figure 3 · p. 9

5 RELATED WORK

Efficient sequence models use recurrent or linear-time token mixers with fixed-size states, while associative memories offer different update rules and structures. This work connects GSA's shared-slot factorization with Oja- and Delta-based correction on the key and value sides.

5 RELATED WORK

6 CONCLUSION

GSA2 applies complementary associative correction to both sides of a shared latent-slot memory, with strong performance across several benchmark families. The paper also states its AI-use, ethics, and reproducibility disclosures.

6 CONCLUSION

  • GSA2 uses Oja Rule-2 for key–slot correction and Delta Rule-2 for slot–value correction, decoupling erase and write control in both stages.Compact WY representations support chunkwise training while recurrent memory remains fixed-size.Source: §6 · p. 10

AI USE STATEMENT

  • Generative AI tools were used only to aid and polish writing, and the authors reviewed the AI-assisted content and take responsibility for the final work.Source: AI USE STATEMENT · p. 10

ETHICS STATEMENT

  • The work uses publicly available training and evaluation datasets and reports no human subjects or sensitive personal data collection.The authors report no additional method-specific ethical concerns beyond those generally associated with large language models.Source: ETHICS STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

  • The paper provides its architecture, recurrent update rules, and main training settings; the appendix supplies detailed derivations and additional experimental protocols.Complete source code is stated to become publicly available upon acceptance.Source: REPRODUCIBILITY STATEMENT · p. 10

A.1 GATED OJA RULE

The Gated Oja Rule first decays the recurrent state and then applies an Oja correction. Its recurrence is expressed through auxiliary memory terms and a WY-style form that supports efficient chunkwise computation.

A.1 GATED OJA RULE

  • The recurrent state is first subjected to diagonal value-side decay, followed by the Oja correction.The appendix assumes strictly positive decay entries so the element-wise divisions are well defined.Source: §A.1 · p. 13
  • The partially expanded recurrence separates pseudo memory decay from pseudo incremental memory.The matrix products are time ordered, with larger time indices on the left.Source: §A.1 · p. 13, §A.2 · p. 14
  • The appendix defines element-wise cumulative decay and prefix- and suffix-decay matrices for the chunkwise representation.The stated state orientation is St ∈ RDv×Dk, with vt and αt in RDv and kt and qt in RDk.Source: §A.1 · p. 13
  • A WY-style representation and UT transform reduce non-matmul FLOPs to support hardware-efficient training computation.The resulting lower-triangular systems are solved by iterative row-wise forward substitution in float32.Source: §A.1 · pp. 13–14

A.2 GATED OJA RULE-2

Gated Oja Rule-2 decouples erase and write control for key-side correction. Its gate-aware recurrences preserve a WY-style structure and use a UT transform for efficient chunkwise computation.

A.2 GATED OJA RULE-2

  • Gated Oja Rule-2 replaces the shared scalar gate with a value-side erase gate and a channel-wise write gate.Under the stated state orientation, the erase gate modulates vt and the write gate modulates kt.Source: §A.2 · p. 14
  • The method retains the WY factorization structure with modified auxiliary recurrences for the decoupled gates.Applying the UT transform to the gate-aware recurrences reduces non-matmul FLOPs and enables efficient chunkwise computation.Source: §A.2 · p. 15
  • The erase gate is absorbed into correction-side factors affecting M[t] and W[t], while the write gate acts only on U[t].Because the gates are separate, no shared Diag(β[t]) factor can be factored from both branches.Source: §A.2 · p. 15

A.3.1 TRAINING DETAILS

The appendix specifies a shared training recipe and describes language-modeling, reasoning, and retrieval evaluations. It also defines a controlled H20 throughput benchmark across sequence lengths.

A.3.1 TRAINING DETAILS

  • Models use approximately 1.3B parameters and train on 100B FineWeb-Edu tokens with the TinyLlama v1.1 tokenizer.Training uses a 4K-token sequence length; hybrid models interleave the recurrent mixer with sliding-window attention using a 2K window.Source: §A.3.1 · p. 15

A.3.2 BENCHMARKS DETAILS

  • Evaluations cover WikiText and LAMBADA perplexity, zero-shot commonsense tasks, RULER retrieval, real-world retrieval, and 11 LongBench tasks.The LongBench average is reported across all 11 tasks; RULER follows its official protocol (Hsieh et al., 2024).Source: §A.3.2 · pp. 15–16

A.3.3 TRAINING THROUGHPUT EVALUATION

  • End-to-end training throughput is benchmarked on one H20 GPU at 2K, 4K, 8K, and 16K sequence lengths with approximately 1.30B-parameter models.Each step includes forward and backward computation and optimizer updates; throughput is measured in Kt/s using the median of three GPU-run means.Source: §A.3.3 · p. 16
Loading 2610.02816v1…