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

2026-10-02

AI Summary

This study addresses the limitation of linear attention models in effectively correcting key-value associations due to their fixed-capacity recurrent memory. To this end, we propose GSA2, an architecture that achieves precise bilateral memory updates for both keys and values via shared latent slots. Methodologically, we introduce the Gated Oja Rule-2 for key-side correction, decoupling it from value-side updates governed by the Delta rule. Furthermore, we design a gated slot attention mechanism alongside a hardware-friendly chunked parallel training algorithm. Experimental results demonstrate that GSA2 significantly outperforms strong baselines across multiple benchmarks while strictly preserving linear time complexity and constant-memory decoding efficiency.

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.

Problem

• Linear Attention
• Recurrent Memory
• Memory Correction
• Delta Rule
• Slot-based Architecture

Innovation

• Linear Attention
• Gated Slot Attention
• Associative Memory Correction
• Gated Oja Rule
• Chunkwise Parallel Training