PaperlayerView exact revision on arXiv

Source-linked AI summary

Gated Slot Attention for Efficient Linear-Time Sequence Modeling

Yu Zhang, Songlin Yang, Ruijie Zhu, Yue Zhang, Leyang Cui, Yiqiao Wang, Bolun Wang, Freda Shi, Bailin Wang, Wei Bi, Peng Zhou, Guohong Fu

arXiv:2409.07146v2cs.CL

TL;DR

Softmax attention poses long-sequence and inference-memory challenges, while linear recurrent models still struggle with recall and training from scratch is costly. GSA adds GLA-inspired gating to ABC in a two-pass linear-attention design. The paper reports strong recall-intensive and T2R performance while retaining a compact state.

  • Problem

    Softmax attention has quadratic complexity and a growing inference cache, while linear recurrent models struggle with in-context retrieval and remain costly to train from scratch.

    Source: §1 · pp. 1–2

  • Method

    GSA enhances ABC with GLA-inspired gating and frames the model as two-pass GLA linked through softmax.

    Source: §1 · p. 2, §7 · p. 10

  • Results

    GSA matches language modeling and understanding performance, outperforms other linear models on recall-intensive tasks, and improves T2R results over recurrent models and other linear-model T2R methods.

    Source: §1 · p. 2

  • Takeaways & Limitations

    GSA combines context-aware reading and forgetting with a small state, improving training and inference efficiency within the evaluated settings.

    Source: §7 · p. 10

    The authors report no long-context results because their pretrained models were relatively small, and GSA underperformed Transformers on Phonebook Lookup and Needle-In-Haystack evaluations.

    Source: §6 · p. 10

Abstract

from arXiv · show

Linear attention Transformers and their gated variants, celebrated for enabling parallel training and efficient recurrent inference, still fall short in recall-intensive tasks compared to traditional Transformers and demand significant resources for training from scratch. This paper introduces Gated Slot Attention (GSA), which enhances Attention with Bounded-memory-Control (ABC) by incorporating a gating mechanism inspired by Gated Linear Attention (GLA). Essentially, GSA comprises a two-layer GLA linked via $\operatorname{softmax}$, utilizing context-aware memory reading and adaptive forgetting to improve memory capacity while maintaining compact recurrent state size. This design greatly enhances both training and inference efficiency through GLA's hardware-efficient training algorithm and reduced state size. Additionally, retaining the $\operatorname{softmax}$ operation is particularly beneficial in "finetuning pretrained Transformers to RNNs" (T2R) settings, reducing the need for extensive training from scratch. Extensive experiments confirm GSA's superior performance in scenarios requiring in-context recall and in T2R settings.

1 Introduction

Linear recurrent models offer efficient inference but remain challenged by recall-intensive tasks, while softmax attention brings quadratic training costs and growing inference caches. GSA combines ABC with gating and a two-pass linear-attention formulation; evaluations report strong recall and T2R results alongside efficient training and inference.

1 Introduction

  • Softmax attention has quadratic complexity, while its inference-time key-value cache grows with generation length and burdens memory and throughput.Linear attention alternatives can be run recurrently with constant memory during inference.Source: §1 · p. 1
  • Linear recurrent models struggle with in-context retrieval, and training them from scratch on trillions of tokens remains prohibitively expensive.T2R uses only a few billion finetuning tokens, but changing from softmax to a linear-attention kernel can create performance discrepancies.Source: §1 · p. 2
  • ABC retains softmax and can use a smaller state than other linear-attention models, but its language-modeling performance and training speed have limited adoption.These properties motivate GSA's effort to improve ABC while preserving its T2R suitability.Source: §1 · p. 2
  • GSA reformulates ABC as two-pass linear attention linked by softmax and adds a gating mechanism inspired by GLA.The reformulation enables use of FLA's hardware-efficient chunkwise implementation.Source: §1 · p. 2
  • GSA matches language modeling and understanding performance, outperforms other linear models on recall-intensive tasks, and improves T2R results while retaining a small state.It has similar training speed to GLA and faster inference due to its smaller state size.Source: §1 · p. 2

2 Background and Preliminary

The background frames attention as key-value memory: standard Transformers accumulate unbounded slots, while ABC bounds the memory and GLA adds input-dependent gating. These approaches trade memory and computation against retrieval and forgetting behavior.

2.1 Transformers as Unbounded Key-Value Memories

  • Causal softmax attention reads from key-value memories, assigning each token a slot and producing outputs from normalized attention scores.The causal mask prevents access to future positions.Source: §2.1 · p. 2

2.2 ABC [63]: Linearizing Attention with Bounded Memory Control

  • ABC bounds memory by fixing the number of slots and allows multiple tokens to write into a slot, unlike sliding-window attention, which discards information outside its window.Its two-pass linear-attention formulation supports chunkwise training; context-aware reading and softmax are proposed to help its compact state retain retrieval capacity.Source: §2.2 · p. 3

2.3 GLA [96]: Linear Attention with Gating Mechanism

  • GLA uses data-dependent decay to selectively retain or forget state, and an outer-product gate parameterization permits chunkwise parallel training.The gate parameters depend on the input; the equivalent recurrence treats the key as an input gate and the decay as a forget gate.Source: §2.3 · p. 4

3 Method

ABC can retain old information indefinitely and favor sentence-initial tokens, limiting memory reuse and recency. GSA adds data-dependent gating and a two-pass GLA formulation, then uses stacked GSA blocks for token and channel mixing.

3.1 Motivation: Issues with ABC

  • ABC lacks forgetting and favors sentence-initial tokens, which impedes memory reuse and conflicts with the recency bias of natural language.The first token writes maximally to every slot, while cumulative softmax can weaken later tokens’ writing strength and challenge long-context retention.Source: §3.1 · p. 4

3.2 Gated Slot Attention (GSA): ABC with gating mechanism

  • GSA adds data-dependent gating to ABC to forget historical information and introduce a recency inductive bias.Each slot uses a scalar gate to blend its previous key and value memories with the current token’s key and value.Source: §3.2 · p. 4
  • GSA can be written as two-pass GLA, enabling adaptation of GLA’s hardware-efficient chunkwise training algorithm.The recurrent representation is shown in Figure 1.Source: §3.2 · p. 4, Figure 1 · p. 5

3.3 Neural Architecture

  • The architecture stacks GSA token-mixing layers with GLU channel-mixing layers and uses multi-head attention.With 64 memory slots and four heads, the parameter count is kept approximately equal to 4d^2.Source: §3.3 · p. 5, Figure 2 · p. 5

4 Experiments

The experiments compare GSA with language-modeling and recall-intensive baselines, test its efficiency, and evaluate finetuning pretrained Transformers to recurrent models. GSA performs strongly on recall-intensive and finetuning tasks while retaining a compact state.

4.1 Language Modeling

  • The language-modeling experiments train 1.3B and 2.7B models on Slimpajama for 100B tokens and compare them with Transformer and subquadratic baselines.The comparison includes Xfmr++, Mamba, RetNet, GLA, and HGRN2.Source: §4.1 · p. 5

4.1.1 Results on commonsense reasoning tasks

  • GSA matches HGRN2 on commonsense tasks with an equally sized hidden state and beats GLA and RetNet with a smaller state.The evaluated zero-shot results cover 1.3B and 2.7B models.Source: §4.1.1 · p. 6, Table 1 · p. 6

4.1.2 Results on in-context recall-intensive tasks

  • MQAR tests retrieval of multiple associative key-value pairs, and its results validate GSA’s effectiveness on this synthetic diagnostic task.The setup uses sequence length 512 and 64 key-value pairs.Source: §4.1.2 · pp. 6–7
  • On recall-intensive real-world tasks, GSA outperforms all subquadratic baselines without a larger state, while Xfmr++ achieves the best average performance.The authors point to context-aware memory readout and gating-based forgetting as possible explanations.Source: §4.1.2 · p. 7

4.1.3 Ablation

  • Ablations find that gating improves language-modeling perplexity, post-pass softmax non-linearities help, and 64 slots balance performance and efficiency.The slot count was selected from the tested alternatives of 32, 64, and 128.Source: §4.1.3 · p. 7, §4.2 · p. 7

4.1.4 Efficiency

  • GSA maintains training throughput comparable to GLA despite two GLA passes and achieves faster inference than RetNet and GLA with its smaller state.Recomputing recurrent hidden states lowers memory consumption but slightly reduces training throughput.Source: §4.1.4 · p. 7, Figure 4 · p. 8

4.2 Finetuning Pretrained Transformers to RNNs

  • Mistral-to-GSA outperforms GLA and RetNet after finetuning, and with 100B tokens exceeds RWKV6 and Mamba, which were trained on over 1T tokens.The authors attribute the finetuning advantage over other approaches to preserving softmax.Source: Table 3 · p. 8, §4.2 · p. 9
  • Across four tasks, GSA consistently outperforms other subquadratic models; Mistral-to-GSA also performs better overall than RWKV6 and Mamba.The cited comparison reports these scratch-trained models used over 1T tokens.Source: §4.2 · p. 9, Table 4 · p. 9

5 Related works

Related work frames linear attention as a route to efficient recurrent inference and hardware-efficient training, while identifying memory capacity and T2R transfer as open concerns. GSA is connected to two-layer fast-weight models and distinguished from approaches that approximate rather than retain softmax.

Related work

  • Mamba’s I/O-aware design reduces data-transfer costs, but its non-matmul recurrence limits tensor-core use and state materialization constrains recurrent state size.Mamba2 [16] instead adopts a linear-attention-like approach to support hardware-efficient training.Source: §5 · p. 9
  • Linear attention’s chunkwise form supports parallel training and recurrent inference, while gating and decay variants seek greater expressiveness without losing matmul-based parallelism.GSA uses two-pass gated linear attention to further enhance capacity while retaining hardware-efficient training.Source: §5 · p. 9
  • Linear attention can be viewed as a single-layer fast-weight network, whereas ABC and GSA form a two-layer fast FFN linked by softmax.The greater capacity of a two-layer FFN could explain why GSA achieves similar or better performance with a smaller state, especially on long and recall-intensive tasks.Source: §5 · pp. 9–10
  • Prior T2R methods approximate attention through learned features, frequency-domain kernelization, or multi-stage distillation, while GSA directly retains softmax [10] [7] [90].Zhang et al. learn a feature map to mimic softmax rather than including the operator in the resulting model.Source: §5 · p. 10

6 Limitations and future work

The evaluation omits long-context results because the pretrained models are relatively small, and copy-oriented tasks show weaknesses against Transformers. The authors identify hybrid approaches, additional model comparisons, and two-pass DeltaNet as future directions.

Limitations and future work

  • Long-context results are not reported because the pretrained models are relatively small compared with models trained on trillions of tokens.The authors view Table 4 as positive indication and suggest larger models and token horizons for future evaluation.Source: §6 · p. 10
  • On Phonebook Lookup and Needle-In-Haystack, GSA performs worse than Transformers on copy-oriented tasks, reflecting limits with precise local token shifts and comparison.The authors suggest hybrid GSA-attention models, Transformer distillation, JRT prompts, or combining with YOCO as possible mitigations.Source: §6 · p. 10
  • The study leaves an analogous two-pass DeltaNet unexplored and calls for comparison with xLSTM, Mamba2, TTT, and Longhorn.Parallel DeltaNet [98] parallelizes DeltaNet’s delta-rule computations over sequence length, improving training efficiency.Source: §6 · p. 10

7 Conclusions

The paper introduces Gated Slot Attention (GSA), which adds GLA-inspired gating to ABC and frames the model as two-pass GLA for hardware-efficient training. Experiments show advantages in recall-intensive tasks and T2R settings.

Conclusions

  • GSA enhances ABC [63] with a gating mechanism inspired by GLA [96].The paper frames GSA as two-pass GLA to use hardware-efficient implementations for training.Source: §7 · p. 10
  • Context-aware memory reading and forgetting increase GSA’s implicit capacity despite its small actual state size, improving training and inference efficiency.Source: §7 · p. 10
  • Experiments demonstrate GSA’s advantages in in-context recall-intensive tasks [4] and finetuning pretrained Transformers to RNNs [41].Source: §7 · p. 10

A Linear Attention and its Chunkwise Form

Linear attention replaces softmax with a decomposable kernel to enable recurrent computation with bounded memory, but its additive update can struggle to forget unrelated information. Chunkwise computation and gating support parallel training while balancing computation and memory costs.

A Linear Attention and its Chunkwise Form

  • Linear attention replaces softmax with a decomposable kernel, and feature mapping applies ϕ to each input.A normalization term similar to softmax can be removed in vanilla linear attention to avoid potential gradient explosions [66].Source: A Linear Attention and its Chunkwise… · p. 18, §A.1 · p. 18
  • Associativity yields a recurrent matrix-valued state updated additively with key-value outer products, reducing computation from O(T^2d) to O(Tmd).This bounded-memory form is more efficient when sequence length T is much greater than m and d.Source: A Linear Attention and its Chunkwise… · p. 18
  • At inference, linear attention maintains O(md) hidden memory instead of an attention KV cache, but its additive update can retain unrelated information.Source: A Linear Attention and its Chunkwise… · p. 18
  • Gating controls information flow for selective memory reading and writing; variants use data-independent scalar decay or input-dependent, finer-grained gates.One formulation couples keys to forget gates to reduce parameters and improve efficiency.Source: A Linear Attention and its Chunkwise… · p. 18

A.1 Hardware-Efficient Training

  • The recurrent form limits training parallelism, while the parallel form remains quadratic because the causal mask prevents reordering computation as KV-first.Source: §A.1 · p. 18
  • Chunkwise computation uses the associative gated recurrence for parallel scans, balancing recurrence and parallelism while reducing non-matmul FLOPs and memory/I/O overhead.Its hardware-efficient implementation can provide practical wall-clock speedup against standard attention.Source: §A.1 · pp. 18–19
  • Cumulative decay factors can be absorbed into Q and K to derive vectorized updates with inter-chunk and intra-chunk terms.The intra-chunk term can be numerically unstable because cumulative products may exceed 1; secondary chunking addresses this issue [96].Source: §A.1 · p. 19

B Details for GSA

GSA is implemented as two GLA passes, with chunkwise procedures for its forward and backward computations. The method uses selective checkpointing to reduce memory needed for backpropagation.

B Details for GSA

  • The first pass absorbs cumulative forget gates into queries and keys, using slot representations I as values.The preprocessing step computes chunkwise cumulative forget gates.Source: B Details for GSA · p. 19
  • First-pass outputs are softmaxed and used as queries in the second GLA pass, which uses I and V as its key and value vectors.The two-pass algorithm is decomposed into inter-chunk recurrence and intra-chunk parallel computation.Source: B Details for GSA · pp. 19–20
  • Gradient checkpointing trades recomputation for lower memory use by saving the inputs and only the first GLA pass output among the forward intermediates.The authors identify saving all intermediates as a memory-management challenge.Source: B Details for GSA · p. 20
  • The backward pass reverses the order of the two GLA backward passes, combines their gradients for I, and decomposes the forget-gate gradient into two parts.The two parts are Q⊙dQ−K⊙dK and O⊙dO−V⊙dV.Source: B Details for GSA · p. 20
  • The forward and backward procedures use chunkwise parallelism, with a secondary chunking strategy in the implementation to ensure numerical stability.The paper provides detailed hardware-efficient forward and backward procedures.Source: B Details for GSA · pp. 19, 20, 22

C.1 Language Modeling

The language-modeling section identifies the compared Transformer and subquadratic architectures, specifies a shared from-scratch training setup, and accounts for recurrent state sizes.

C.1 Language Modeling

  • The comparisons include Xfmr++ and the subquadratic architectures Mamba, RetNet, GLA, and HGRN2.Xfmr++ is described as Llama-like; the listed alternatives include state-space and linear-attention models.Source: §C.1 · pp. 22–23
  • All models are trained from scratch with the same recipes on a 100B-token Slimpajama subset, using 64 memory slots.The experiments use the Mistral tokenizer and the FLA library; training runs use 32 Nvidia H800 GPUs.Source: §C.1 · p. 23
  • Reported total recurrent state sizes are 64Ld for Mamba, 128Ld for GSA and HGRN2, 256Ld for GLA, and 512Ld for RetNet.These are the state-size accounting figures given for the compared architectures.Source: §C.1 · p. 23
Loading 2409.07146v2…