Gated slot attention improves memory in efficient linear models
Gated Slot Attention-2: Two-Sided Associative Memory Correction in Linear Attention
Machine Learning
Summary
Managing memory in fast, linear attention models is tricky because they keep a limited fixed memory that updates over time. The authors found two separate memory correction techniques that work well together: one that tweaks memory related to 'keys' and one that updates 'values' associated with those keys. They created a new combined method, called Gated Slot Attention-2, that improves how these memories are corrected and stored. This method keeps the model efficient while performing better on various sequence processing tasks.
What this means in practice
- •For machine learning engineers: Enhance efficient sequence models for natural language or time-series data by improving fixed-size recurrent memory correction.
- •For hardware architects: Design efficient parallel training hardware using the derived chunkwise algorithm for linear attention models with improved memory correction.
Authors
Ruijie Li, Shengnan Ding, Weimin Zhang, Derick Tang, Zhanpeng Zeng, Qinsong Zeng, Ming Chen, Jiaxi Hu, Yuxuan Liang
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.