Gated Slot Attention for Efficient Linear-Time Sequence Modeling
Authors: Yu Zhang
Songlin Yang
Rui-Jie Zhu
Yue Zhang
Leyang Cui
Yiqiao Wang
Bolun Wang
Freda Shi
Bailin Wang
Wei Bi
Peng Zhou
Guohong Fu
Venue: NeurIPS 2024 poster (2024)
Keywords: sequence modeling
language modeling
hardware-efficient
linear attention
Primary area: natural_language_processing

Abstract: 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.

Gated Slot Attention for Efficient Linear-Time Sequence Modeling
Yu Zhang' Songlin Yang?* Ruijie Zhu
Yue Zhang' Bailin Wang*
Leyang Cui
Yiqiao Wang² Bolun Wang Freda Shi
Wei Bi
Peng Zhou
Guohong Fulf
1 School of Computer Science and Technology, Soochow University, China
Massachusetts Institute of Technology University of California, Santa Cruz
Tencent AI Lab LuxiTecc
University of Waterloo
soraann.cs@ooutlooo. com
sangs166@mit.eem
https://github http://uugggnngacccccoooooooo
Abstract
Linear attention Transformers and their gated variants, celebrated for enabling parallel training and efficient recurrent inference, still fall short in recall-intensive tasks compared traditional Transformers and demand significant resources for training from scratch. This paper introduces Gated Slot Attention (GSA), which enhances Attention with mounded-memory-Coe gating mechanism inspired by Gated Linear Attention (GLA (96).. Essentially, GSA comprises two-layer GLA linked via softmax, utilizing context-aware mem- ory reading and adaptive forgetting to improve memory capacity while maintaining compact recurrent state size. This design greatly enhances both training and infer- ence efficiency through GLA's hardware-efficient training algorithm and reduced state size. Additionally, retaining the soft operation particularly beneficial in "finetuning pretrained Transformers to RNNs" (T2R 41]) 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.
(ABC (63 by incorporating
Introduction
Transformers
88 have emerged as the predominant architecture for most, if not all, sequence modeling tasks. Nevertheless, the quadratic complexity of softmax-based standard attention (SA) poses significant challenges for long sequence modeling (c.g.. video understanding and biological sequence modeling). In the context language modeling, where sequence lengths are moderate. training efficiency is generally not primary concern. However, during inference, the Key-Value (KV) cache [34 64 grows linearly with the generation length, resulting substantial memory burdens and throughput bottlenecks due to high I/O costs.
Linear (kernelized) attention [42 and its gated variants 96
have received interest max attention. These models demonstrate strong performance in
82
68
61
16 69
as promising alternatives to sof language modeling and understanding tasks. Notably, they be reframed as RNNs during inference, achieving constant memory complexity and thereby significantly enhancing inference efficiency.
However, two key issues persist with these models: (i) Performance- wise, recent research indicates that linear recurrent models still struggle with tasks requiring in-context retrieval or learning [ 37
"Equal contributions. Work conducted during Yu Zhang's internship Tencent Al Lab. Conresoonding authors.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

From this perspective, Transformers are equipped with unbounded number of memory slots, which grow linearly with respect to the sequence length [57] (i.e., m = t for step f)-a new key k; W, _ is assigned with unique memory slot upon its introduction. This leads to simple similar way. 28 and there is a findamental recall-memory trade-off [3 19 where all inference-time time-constant- memory models face inherent limitations. (ii) In terms of training efficiency, while linear attention supports hardware-efficien chunkwise training [96] as implemented in internttttttooo (FLAA
(95)) training from scratch on trillions of tokens remains prohibitively expensive. A paradigm, "finetuning pretrained Transformers to RNNs" (short for T2R [41). has recently gained great attention [01 [4 (3 7 [0]. This approach circumvents the high cost of training from scratch by requiring
only a few billion tokens for finetuning- -bbuut 1-3% of the total cost. However, linear attention uses different kernel method from softmax, leading to performance discrepancies when finetuning pretrained softmax attention models to linear attention [0O1.
To address these issues, we revisit the Attention with Bounded- Memeory Control (ABC) model [63] which retains the softmax operation, thereby reducing training- fintiininnnnnuiigggggg discrepancies between standard and linear attention, making it ideal for T2R settings. Additionally, ABC enables more effective state utilization, requiring less state size to achieve similar performance, as observed in Peng al. 633 This results in more efficient inference and potentially expands the Pareto frontier of the recall-memory tradeoff [3]. However, ABC has not gained significant attention due to its mediocre language modeling performance and slow training speed.
In this work, we first reformulate ABC as two-pass linear attention linked via softmax, allowing us to leverage the bardware-efficient chunkwise implementation from FLA [95] for more efficient training. We then identify several limitations of ABC and propose new model, dubbed Gated Slot Attention (GSA), which is essentially a gated version of ABC, following the recent trend of enhancing linear
attention with gating mechanisms 99 69
61
16
6 62
65
Our extensive evaluation shows that GSA not only matches performance in language modeling and understanding tasks but also significantly outperforms other linear models in in-context recall- intensive tasks 3 4), without requiring a large state size like RetNet [82] or GLA [96). In the T2R finetuning setting, we found that finetuning Mistral-7B (39] to GSA A surpasses large recurrent language models (e.g.. RWKV6-7B, Mamba-7B) and also outperforms finetuning Mistral-7B to other linear models (e.g.. RetNet, GLA) and other T2R methods like SUPRA [54], verifying the importance of retaining the softmax operator. Finally, we remark that GSA achieves similar training speeds to GLA while offering an inference speedup due to its smaller state size.
Background and Preliminary
2.1 Transformers as Unbounded Key-Value Memories
Given X = ==x ærl _ E 7xx, where T is the sequence length and x; E Rd is the i-th input vector with d dimensions, SA with causal masking computes the output matrix:
O = [[(QK] ② M)V,
where Q, K, V Txd are linear mappings of the input X via learnable weights Wg: WR. W. E Ddxd M [M = 1 > j o.w. oc) is the causal mask to prevent future information leakage, C denotes clement-wise production, and f(-) is softmax (-).
Generally, K, V can be viewed as neural key-value memories K,, V, E mxxd respectively 81 24 where m is the number of memory slots. At step I, the query q, = W E Rd first attends to the key memories K, to retrieve relevant information, which is then summarized into og by computing a weighted sum of the value memories V, [104) where the weights are the normalized attention scores:
O, V. ((K.g))
memory updating rule: K,- = K, This mechanism, however, comes at the cost of quadratic time complexity in terms of the sequence length for training and (TTd) time/memory complexity for inference (64), posing challenges for large-scale models.
1 U [k.). The value memories V, are updated

---

2.2 ABC [63 Linearizing Attention with Bounded Memory Control
From a key-value memory perspective, the training and inference complexity of self-attention (SA) can be reduced by fixing the number of memory slots to constant size < T [2 51 63). One straightforward way to achieve this is by employing a first-in-first-oul memory management strategy, commonly known as sliding window attention (SWA). However, SWA is inefficient because it discards all information outside the window, leading to poor performance in balancing the recall-memory tradcoff [3]. To achieve acceptable performance, SWA often requires a large window size (c.g.. 4,096 tokens in Mistral 39)), which diminishes its advantage over to global attention.
When the number of tokens in a sequence exceeds the number of memory slots, it becomes necessary store information from multiple tokens single slot. To address this challenge, Peng et al. [63] propose the Attention- with Bounded-memory- Cocnool (ABC) mechanism, which allows multiple tokens to be written into a single slot:
K, K to. Sk E Rmxd V Vi=V-++0,V E E
O, f(K qi) R (3)
where
a; exp ( (w
Rm E
ON Σ
E (0, 1)
Here, (Φ:), represents the writing intensity of the ith token to the jth slot, obtained using cumulative softmax function (ef. 63 footnote 5)), which can computed with prefix sum.
ABC two-pass linear attention. bears a resemblance to linear attention
The outer nuter-product-buee based additive memory update rule in Eq. 42 which involves the following recurrence
S. = S k @ ve
R dxd
O, = S 9: E R
(5)
We denote this linear attention operator that computes o, from q; k; and Us (Eq. 5 by {o,)T T.A ((g, k, v)))= We show that the ABC operations can be written as
(o = LA((g,, ks, 1.)T= (o.) = TA(((sftmmax(()) (0)) & m. o.)
where 0f E RT O E Rd. Therefore, ABC can enjoy hardware-efficient linear-ting chunkwise training (96), implemented the FLA library (95).
Remarks on state size. Peng et al. 663 empirically demonstrated that ABC requires smaller state size to achieve comparable performance to other linear attention models, resulting in improved inference efficiency. We offer the following intuitive explanation: the new query of aggregates the entire history through the initial pass of linear attention, making it more context-awand aware and better at locating desired items for retrieval. The subsequent softmax operator helps mitigate the attention dilution issue (66). From the perspective of Hopfield networks, softmax can exponentially increase the memory size [45] Together, these factors suggest that ABC may possess implicit large memory capacity, even with a small actual recurrent state size.
2.3 GLA 96 Linear Attention with Gating Mechanism
Linear attentions underperform softmax-attention Transformers in language modeling by notable margin. RetNet (82) and TransnormerI_LM (68 incorporate a data-independent exponential decay factor for memory update
= 7 + k, & 0, E BBdxd
where y E (0, 1) is scalar data- lata-independent decaying factor; that is, the decay rate is fixed across time steps and hidden channels (under the same head), disrespect the input tokens. RetNet has shown better language modeling performance compared to vanilla linear attentions thanks to the decaying mechanism.
For simplicity, omit the normalization term, which has been shown be unnecessary
66
52
2
96

---

However, research in recurrent neural networks (RNNs) has shown that data- deta-dependent decay (or forget gates) crucial for selectively retaining and forgetting information [2] 26]. thus better
leveraging the fixed recurrent hidden state. This selective mechanism has been revisited in recent
state-space models [29 [6]. Inspired by LSTMss, Gated Linear Attention (GL.A) [52 96 introduces
data-dependent decay parameters G; E (0, 1 yddd to gate the hidden state as follows,
S G2 o S + k2 @ E Rdxd O2 = S q: E R.
[6 show that if gates are parameterized in an outer product form G, = a & B, and a B, E [0, 1jd depend solely on input æz, such recurrence
be rewritten matrix multiplication, allowing for hardware-efficien training with a chunkwise parallel form. In what follows, we will use the following notation GLA (fg.A(fg.))) k v, O2 (=)= for y.= denote this computation. It is common set 3 = 1 as in 196 69 61]. which is also often written in the following equivalent form:
S, Diag (a,gg(,,)SS-111 + k @ Vt-.
Here & can be viewed as the input gate, and 0 can be viewed the forget gate. In gated RNN literature, it is common to couple these two gates via k2 = 1 ar [12 [06 67 In particular, Qin al. [69] proposed HGRN2, which uses this strategy as an improved parameterization of GLA, showing better performance in language modeling.
3 Method
3.1 Motivation: Issues with ABC
We identify two primary limitations ABC's memory update rule. Firstly, lacks a forgetting mechanism, resulting in indefinite retention of items once written into memory slots. This prevents efficient memory reuse by impeding the prompt clearance of slots for new information.
Secondly, the rule introduces an unwarranted inductive bias favoring tokens at the sentence's begin- ning. This contradicts the recency bias natural language, where more recent information is often more relevant. Prioritizing initial tokens over the recent ones conflicts with this inherent tendency in natural language processing.
Specifically, for the first token, the writing strength to all slots is maximized (i.e., Φ1 = 1 E IR"), causing every memory slot retain copy of the first token's 's representation. The absence of forgetting mechanism exacerbates this issue. For subsequent tokens, the writing strength diminishes due the influence of earlier tokens, as result of the cumulative softmnx in Eq. [ This makes challenging for the model to retain later tokens without learning significantly large og, potentially leading to instability in long-context settings, observed by Zhang al. (T00.
3.2 Gated Slot Attention (GSA): ABC with gating mechanism
To address these limitations, we propose Gated Slot Attention (GSA), which incorporates a gat- ing mechanism to simultaneously resolve both issues by: (i) enabling the forgetting of historical information, and (ii) introducing a recency inductive bias, as detailed below.
For each memory slot, the update rule is a simple gated RNN with a scalar data-dependent gating value ar [0, 1].
(K,) og (K,-1); + (1 ax)k: E Rd, (V.), as V -1-)) + oz) , Rd
and these can be written matrix form, which reminiscent HGRN2 169]
K, Diag(a:) K + (1 a:)) @ k. E mxx
V, = Diag(a.) V + (1- 11 ax) v E mmxd
(6)
0, vT softmax K. q:) E Rd
GSA as two-pass GLA. shown below:
is straightforward to see that we can write GSA two-pass GLA as
{o]) GLA (fg,; a4, 04, 1}
{o:)²'
1 = GLA ((softmax (ot), 1-11--0 Cce, Ve, 1, ax)T
Therefore, we can adapt GL.A's hardware-efficien chunkwise training algorithm for GSA training, as shown in $ A and B We illustrate the recurrent representation of GSA in Figure])

---

3.3
Neural Architecture
Outputs
Lincar
Norm
Vi-1
GLU
Normn
Norm
Nx
Gated Slot Attention
GSA
Linear
Ki-1
Inputs
Figure 1: The recurrent representation of GSA. -W means taking at, input.
Figure 2: The backbone of our proposed GSA models.
The overall architecture of our proposed model, GSA, shown in Figure Following the Llama architecture [86], we use stack of L GSA blocks, each comprising GSA token mixing layer followed by a Gated Linear Unit (GLU) channel mixing layer 19 33].
We utilize the multi-head attention mechanism 88 to capture different aspects of the input. For each head h, the input to GSA token mixing is defined as
d. k e <(W" o (W a,), o (W.:.:)
(8)
where @ is the Swish activation following [68]. The forget gate is obtained by a linear transformation followed by sigmoid activation a with a damping factor t [96 [3] af o (Wh,,,)
where the damping factor is to regulate the forget gate value to one, which has been shown to be crucial for long-term dependency modeling (30 667. We feed them into GSA layer to obtain outputs as
described in Eq. 7
1o Je= GSA((d) kR of o D
Finally, we obtain output via
y W. (RMSNorm (Swish (Concat (o)
of ()))
The total number of parameters for W We, W o. and Wo is already 4af, which is the same in single standard softmax-attention layer. To control the overall parameter count, we aim to keep the parameters for W which amount to dHH, relatively small. In practice, balance between efficiency and effectiveness ($ (.)4...). One way to further manage the total parameter count by reducing the number of heads. In practice set H 4, ensuring that Hm < d. This keeps the total number of parameters approximately equal to 4d
set 64 to achieve
Experiments
4.1 Language Modeling
We perform moderate-scale language modeling experiments with 1.3B and 2.7B parameters on Slimpajama corpus [79] for 100B tokens each.
We compare the performance of GSA against Llama Transformer architecture (i.e., Xfmr++ 86 and recent subquadratic architectures including: Mamba (29), RetNet [82), GLA [96] and HGRN2 69
We refer readers for more details on baselines and other experimental setups.
In practice we
For instance, in 1.3B model with H x = 64 4 - 256 and d' = 2, 048, the total number of parameters amount to 4. 125dl introducing only 0. 125d2 overhead.
5

---

4.1.1 Results on commonsense reasoning tasks
Following [29 96], we report the perplexities and zero-shot performance of commonsense reasoning tasks including ARC, & ARC, (ARC-easy. ARC-challenge) [14; Hella. (Hellaswag) [99]. Lamb. (Lambada) [59]. PIQA [8). Wiki. (Wikitex() [55]. and Wino. (Winograde) [73. We note that these tasks are typically short length do not require in-context learning capabilities, thus they do not adequately reflect long-context modeling or in-context learning retrieval abilities. Nevertheless, as shown Table we found that GSA performs comparably to the recent strong model HGRN2 with an equally sized hidden state, while outperforming GLA and RetNet even with smaller state size.
Table 1: The zero-shot results of 1.3B and 2.7B models evaluated by In--valluttionhhrrees 21 denotes number of layer while d denotes the model dimension.
Lamb.
Wiki.
ARC. acc
ARC. acCn
Hella.
Lamb,
PIQA
Wino. acc
State size
Avg.
ppl
ppl
acc
1.3B parumeters with 100B training tokens, L=24,
d= 048
Xfmr+
N/A
15.3
17.1
54.1
27.1
49.3
47.0
70.3
54.9
50.5
Mamba
64 Ld
15.4
17.3
57.1
28.2
50.3
44.4
71.8
52.3
50.7
RetNet
512 T.d
15.4
17.3
57.4
27.9
50.3
44.6
71.7
51.8
50.6
GLA
256 L.d
15.4
17.6
55.4
27.7
49.0
46.4
69.9
54.0
50.4
HGRN2
128 Ld
11.8
16.9
58.1
28.1
51.8
49.4
71.4
52.3
51.9
GSA
Ld
12.6
16.7
58.1
28.2
51.0
47.4
53.4
51.7
2.7B parameters with 100B training tokens, L=32,
d=2,560
Xfmr++
N/A
10.7
15.2
59.8
27.5
54.2
$2.3
72.7
56.2
53.8
Mamba
64 Ld
13.6
15.9
60.7
29.8
53.9
46.4
72.8
53.9
52.9
RetNet
512 Ld
11.9
15.8
59.6
28.1
54.0
49.6
72.3
53.8
52.9
GLA
256 T.d
12.4
15.5
59.2
29.9
54.0
50.4
71.7
55.7
53.5
HGRN2
128 L.d
8.8
14.6
60.8
30.3
58.7
55.4
73.0
54.2
55.4
GSA
128 Ld
9.8
14.8
61.9
30.7
57.0
52.7
73.5
56.0
55.3
4.1.2
Results on in-context recall-intense tasks
While subquadratic models can achieve comparable performance to (softmax- Transformers in language modeling and understanding tasks, their performance on recall-intensive tasks significantly lags behind Transformers and varies greatly across different subquadratic models, observed in many recent studies 3 96 97 Therefore, it is crucial to improve linear models on in-context recall-intensive tasks.
(b) Results the recall- -eaall-nteensivee tasks used in input maximum of 2K tokens.
4 We truncate the
75
State size FDA SWDE SQUAD NQ TriviaQA Drop Avg.
Accuracy
1.3B params / 100B tokens, L=24,
d=2048
50
Xfmr++
N/A 46.0
29.2
41.0
24.8
58.8
21.3
36.9
Mamba
Mamba
Ld 13.9
25.4
18.5
53.5
21.7
27.7
GLA
RetNet
512 x Ld 21.2
27.2
34,0
15.5
52.7
20.0
28.4
RetNet
HORN2
GLA
256 x 26.7
30.6
34.8
21.5
19.1
128
2.56
512
HGRN2 128 x Ld 9.9
23.1
32.0
16.4
55.2
19.1
25.9
Model dimension
GSA
128 x Ld 23.6
29.8
36.0
23.2
57.0
20.9
31.8
(a) Results on the synthetic MQAR task. We adopt the most challeng- ing settings in [2]. utilizing
2.7B params / 100B tokens, 1.=32,
d=2560
Xfmr++
N/A 62.3
30.9
44.3
29.3
61.8
21.4
41.7
Mamba
64 x Ld 21.5
26.7
34.2
21.2
57.0
22.2
30.5
quence length of 512 64 key- value pairs. Xfmr++ with standard at- tention achieves near-perfect results in this settings and is thus omitted for brevity.
RetNet
512 x Ld 24.1
26.1
20.4
57.3
21.8
GLA
256 x Ld 30.3
35.5
36.8
23.3
58.2
21.8
34.3
HGRN2 128 x 15.0
29.9
35.1
20.0
GSA
128 x Ld 39.1
33.5
39.0
26.9
60.8
19.9
36.5

---

MQAR. We first present the results on the multi-query associative recall (MQAR) task [2], diagnostic synthetic task that requires models to retrieve multiple associative key-value pairs from the context. This task has been shown to strongly correlate with language modeling performance [ The results Table 3a validate the effectiveness of GSA.
Real-world tasks. Next, evaluate the zero-shot in- context learning performance on incteliintensive tasks, as assess infor- mation retrieval on FDA 93 and SWDE [49]. which are designed evaluate retrieval from in-context passages scraped from HTML/PDFs. We also evaluate question answering on SQuAD 770,, NQ [46), TriviaQA [40], and Drop (20), where models must ground their answers in in-context documents.
Table 2: Ablation study results for 340M models trained 10B Slimpajama tokens.
used in Arora et al. [
Specifically,
PPI. (
GSA w/ 64 slots
13.51
Ablations on gating mechanism
w/o decay (i.e., ABC)
16.94
w/ data independent decay
15.83
As shown in Table Xfmr++ achieves the best average performance, expected. Meanwhile, GSA outperforms all other subquadratic baseline models by notable mar-
Ablations on non-linearity
softmax
14.03
softmax Swish
13.71
gin without requiring a larger state size. We believe this
softmax + ReLU
13.69
advantage stems from GSA's context-aware memory read-
softmax ReLU
13.95
out mechanism (as discussed in $2.2) and its forgetting mechanism (i.e., the gating mechanism), enabling it to
Ablations slot size
w/ 32 slots
13.74
manipulate finite-sized memory more effectively.
w/ 128 slots
13.46
4.1.3 Ablation
Table 2 presents the results of our ablation studies. Our findings indicate that: (i) the inclusion of the gating mechanism in GSA crucial for improving language modeling perplexity; (ii) applying softmax non-linearities after the first recurrent pass is beneficial; and (iii) using 64 slots strikes an optimal balance between performance and efficiency.
4.1.4
Efficiency
Fig. 4a illustrates the training throughput for four models on single H800 GP To optimize memory usage, we employ the technique of recomputing the recurrent hidden state during the backward pass, as done FLA (95] and Mamba2 [16]. This approach results in reduced memory
consumption (Fig. (4b) at the cost of slightly lower training throughputs (Fig. 4a).
Despite requiring two GLA passes, GSA maintains comparable training throughputs to GLA due to its reduced state size. Since inference is primarily memory-bound inference speed highly correlates with state size. As result, GSA, with its smaller state size compared to RetNet and GLAA, achieves faster inference speeds, as shown in Figure 4c
4.2 Finetuning Pretrained Transformers to RNNs
The concept of finetuning pretrained Transformers linear Transformers for recurrent inference was first introduced in T2R [4]]. This approach uses pretrained language model weights to initialize all parameters, leveraging the similarity between linear attention and softmax attention, and finetunes all parameters, significantly reducing the total training time compared to training from scratch. Kasai al. [41 also introduced parametric feature map, implemented as a learnable MLP layer followed by ReLU, applied after the query/key projections. SUPRA, follow-up T2R, found that the original T2R approach did not perform well the era LLMs, and highlighted the importance of
SSince our pretrained models are neither instruction-tuned nor instruction-aligned, following Arora al. 4 their Cloze Completion Formatting prompts for evaluation. It is noteworthy that results for certain tasks may differ significantly from those obtained using templates.
due to variations in prompt
'Empiricall,, found that 32, 64, and 128 slots result in training throughputs of 46.7K, 44.1K. and 37.1K tokens/s, respectively, under the settings described in the next section. Given the marginal improvement when increasing the slot size from 64 to 128, along with the significant slowdown in training, we chose 64 slots.
We utilize the training throughput benchmark scripts provided by FLA 95 for measurements.
7

---

C Ximme
B
Mamba
GLA
339
Xint++
GSA
CSA mooo
Mamba
45k
37.4
200
RetNet
GLAA
Tolenss
Lakawy (4)
150
GSA
0/4
100
38.0
4SS
15k
Velo Romp
0
50
41.4
2K
x-
16K
20
21
4K
8K
16K
Contest length Bunch size
Memoriis (GiB))
Decoding
(b)
(c)
Figure 4: (a) Training throughput of various 1.3B models single H800 GPU, with fixed batch size containing 16K tokens. "GSA w/o recomp." indicates the use the GSA kernel without hidden state recomputation during the backward pass. Memory footprint (in GiB) of each 1.3B model during training with a batch size containing 16K tokens. [c)]nnerenne] latency (in seconds) of each 1.3B model on single H800 GPU with 2K prefix tokens and batch size of 1.
Table 3: Performance comparison across various 7B models. denotes models using softmax-attention. denotes our results.
Size Tokens
ARC,
ARC.
Hella.
PIQA
Wino.
NQ
TriviaQA
BBH
MMLU
Avg.
Shot(s)
0
3
5
Models trained from scratch (for reference)
RWKV6
7B
1.4T
73.6
44.0
75.2
78.4
68.5
20.9
59.5
23.4
43,9
54.1
Mamba
7B
1.2T
77.6
46.8
77.8
81.0
72.3
25.4
66.2
21.5
33.2
55.7
Llama2
7B
2T
76.4
46.2
76.0
78.0
69.2
26.0
64.2
39.1
45.5
57.8
Gemma
7B
53.2
80.5
79.8
74.0
58.9
63.2
Mistral
7B
?
80.8
54.0
81.1
80.6
74.0
29.7
70.3
56.5
62.4
65.5
Models finetuned from Mistral 7B
SUPRA
7B
+20B
74.6
42.3
74.8
80.1
67.4
28.0
RetNet
7B
+20B
73.3
39.9
72.9
77.8
66.1
16.2
43.0
8.7
26.1
47.1
GLA
7B
+20B
74.6
44.0
75.9
79.2
69.5
22.2
57.8
20.8
28.4
52.5
GSA
7B
+20B
75.9
43.9
76.5
78.7
70.1
23.4
60.7
23,5
32.4
53.9
SUPRA
7B
+100B
76.0
45.7
77.1
79.9
70.3
24.7
60.4
19.8
34.1
54.2
GSA
7B
+100B
76.0
46.9
77.9
78.9
72.6
26.9
65.8
29.3
38.1
56.9
output normalization and a decay mechanism adopted from RetNet 82 performance. As result, SUPRA essentially combines T2R and RetNet by finetuning pretrained Transformers into a RetNet architecture, though excludes Swish output gate.
as critical for finetuning
Settings. In our preliminary experiments, found that the learnable MLP layer unnecessary and could merged into the query and key projections, similar to the approach in Peng et al. (63] We finetuned the pretrained Transformer Mistral 7B [39 to RetNet, well
GLA and GSA models. Following SUPRA, we add ReLU as the feature map activation for RetNet and GLA, which originally used identity feature map without activation ?] and also excluded the Swish output gate. For RetNet, there were no additional parameters; for GLAA the low-rank forget gate, and for GSA, the W., matrix trainable parameters, though both are small in parameter count and negligible in terms of the total model size. We set the peak learning rate to 3 x 10-5 with IK steps of linear warmup following SUPRA. The training length was set 2K tokens, with batch size of 2M tokens. For convenience, we trained on the SlimPajama corpus, while SUPRA used RefineWeb [60] higher-quality corpus. We leave the use of RefineWel for future work.
HHowever, SUPRA reported poor performance with this strategy due to significant discrepancy between training and finetuning, where identity map can lead to negative attention scores, pattern unseen in pretrained Transformers due the nonnegativity of softmax.

---

Main results. Following Jiang et al. 33
Touvron et al. 87]. we evaluated the models on common- Hellaswag [99]. PIQA 8 and Winogrande [73] world
sense reasoning tasks: ARC, and ARC.
knowledge tasks: NQ 46 and TriviaQA 40 and popular aggregated benchmarks: MMLU [I
We observed a clear advantage in finetuning Mistral to GSA compared to GLA or RetNet, confirming our intuition that preserving softmax is beneficial in T2R settings. When trained with 100B tokens, Mistral-to-GS.. outperforms RWKV6 and Mamba on average, even though those models were trained on over IT tokens, thereby reducing the required training data size.
and BBH [85] Results are shown Table 3
Long-context ability evaluation. Following Xiong et al. [94]. we evaluated the models on
Table 4: Long-context performance comparison.
long-sequence tasks, including Qasper [8]. Nar-
Qasper NarrativeQA QUALITY QMSum
rativeQA 14,, QuALITY [58], and QMSum
Models trained from scratch (for reference)
[105]. For each task, the input was truncated to
RWKV6
9.2
14.4
30.8
1.1
16K tokens, which is 8x the training length.
Mamba
5.6
27.9
27.5
0.8
The results are shown in Table Notably, GSA consistently outperforms other subquadratic models across all four tasks. We attribute this to the same factors observed in in-context recall- intensive task settings. Interestingly, Mistral- to-GSA also demonstrates overall better perfor- mance compared to RWKV6 and Mamba, which were trained from scratch on >TT token.
Mistral
25.8
25.1
38.0
5.0
Models finetnne from Mistral 7B 20B tokens
RetNet
11.1
0.0
26.2
0.0
GLA
18.4
17.2
30.9
9.0
GSA
18.8
19.2
32.0
10.0
5 Related works
Matrix-valued linear RNNs with hardware-efficient training. Traditional RNNs (e.g. LSTM
[32]. GRU (12)) maintain 1-dimensional hidden states, which are often too small to capture sufficient information. Recent work emphasizes the importance of expanding the size of recurrent states
61 16]. However, naive state expansion dramatically increases FLOPs and I/O costs, making training impractical. To address this, Mamba introduces an I/O-aware approach, reducing 1/O costs by materializing parameters and hidden states only on SRAM (instead of HBM). However, Mamba's s recurrence cannot expressed matmul form, leading to two key issues: (i) high FLOP count cannot be optimized via tensor cores (the GPU's fast matmul unit), resulting in slower runtimes; and (ii) the recurrent hidden states cannot be compactly represented and must be materialized on SRAM during backpropagation, limiting the recurrent state size due to SRAM constraints.
29
69
96
82
Mamba2 [16] addresses these limitations by adopting linear attention [42 approach that enables sarddire-efffccentt training. Linear attention expands the state using outer products, allowing for both parallel attention-ste -tyle computation and recurrent inference (also known as state-space duality in Mamba2). The chunkwise algorithm interpolates between parallel and recurrent forms, enabling hardware-efficien linear-time training [33
%6 However, vanilla linear attention underperforms softmax attention in various tasks. Recent research has explored incorporating various decay or gating mechanisms to enhance model expressiveness and performance while maintaining matmul-based parallelism and chunkwise training. These include head-wise data-independent decay 82 68 head-wind wise data-dependent decay [6) 6 [3); and channel-wise data-dependent decay 996
12 43 69 61). GSA leverages two-pass gated linear attention to further enhance capacity while
allowing hardware-efficicn training.
Fast weight RNNs.
Fast weight programming 77), a classical concept intensively investigated in
52] has been shown to be closely related to (linear) The core idea involves using slow network to produce rapid context-dependent weight modifications for the fast network. linear attention, the fast network a single-layer FFN with weight matrix S, (Eq. 5). while the slow networks are the query/key/value projections.
deep learning 55
103
74
76
56
75
35
36
Transformers 75
Linear attention is known to suffer from limited memory capacity (75), potentially due to the constraints of single-layer FFN without a large representation. In contrast, ABC and GSA A can be viewed as implementing a two-layer fast FFN with either additive update rule or gated update rule where the weight matrices are K, and V, connected by the softmax activation function
74.
52
(Eq. 3 and 6). This structure resembles DeltaMLP [35,, which uses delta update rule [92 75 98

---

and multi-layer (potentially beyond two layers) fast FFN. The greater capacity of a two-layer FFN compared to a similarly sized single-layer FFN could explain why GSA requires a smaller state size to achieve similar or even better performance, especially in long sequence and recall-intensive tasks.
Finetuning Transformers to RNNs. As discussed, this paradigm could significantly reduce the training cost for large-scale recurrent language models. The idea of distilling Transformers to RNNs to improve inference efficiency can be traced back to Gerstenberger al. [23]. In the following, we briefly introduce some recent works that complement those already mentioned in $4.2 Zhang al. [01] highlight the desirable properties of softmax, such attention spikiness and dot-product monotonicity, and employ a learnable MLP layer to approximate softmax behavior using logit distillation loss (while freezing other parameters). Chen et al. [10] introduce DiJiang, an effective method for approximating attention distributions using the Discrete Cosine Transform (DCT) to enable rrequency-domaii kernelization, leading faster feature mapping. Bick al. 7 propose multi-stage distillation approach, aligning attention distributions (similar to Hedgehog [O01), hidden states, and output logits to transfer knowledge from a pretrained Transformer teacher student Mamba model. Wang et al. [90] distill Transformer-biie based LLMs into hybrid Hlmba-Attentin architectures in the spirit of Ren al. 71. Lieber al. 447, Waleffe al. [8]. However, they freeze the FFN weights, while Choi 13 suggest that it might be more effective to unfreeze them. In this work, we highlight the importance of the softmax operator, as discussed in Zhang et [T0I, except that GSA directly incorporates softmax, while Zhang et al. [TOI] learns feature map to mimic softmax, without actually including any softmax operator in the resulting model.
Limitations and future work
Due to the relatively small scale of our pretrained models (compared to large-scale models trained on trillions of tokens), we did not report any results on long-contex: tasks, as the performance Tableepoovides provides positive indications of GSA's long-context capabilities, and training on a larger token horizon and with larger models would address this. For copy-oriented tasks, we observed negative results the Phonebook Lookup (38) and Needle-In- Haystack evaluations compared to Transformers, revealing the fundamental limitations of linear recurrent models in handling "precise local token shifts and comparison", as discussed in Arora et al. 3. Nonetheless, we expect this limitation could be significantly mitigated by pretraining a hybrid GSA-attention model, as recently explored 3 71 89
would all be poor. However, we believe Table
47 98, or by distilling pretrained Transformers into hybrid GSA-attention models, as in Wang et al. [90 or using different training objectives with
JRT prompts, as in Arora et al. or combining with YOCO 83 25
GSA follows GLA in using a gated update rule, although we acknowledge recent work on Parallel DeltaNet (98) which parallelizes the delta update rule computations in DeltaNet (75) over sequence length, significantly enhancing training efficiency. The delta rule is known to improve in-context retrieval ability 775 198. aligning with one of the objectives of this work. We did not explore the analogous two-pass DeltaNet, but we leave this for future investigation, which would bring the approach closer to the original DeltaMLP (35), discussed earlier. It would also be beneficial to compare GSA with more recent strong RNN models, such as xLSTM 6), Mamba2 [6]. TTT
and Longhorn 48
7 Conclusions
This work introduces Gated Slot Attention (GSA), which enhances ABC [3] with gating mech- anism inspired by Gated Linear Attention (GLA (96) ). By framing GSA as a two-pass GLA, we can leverage hardware-effice efficieetf implementations of GLA 95 to train GSA. As such, GSA benefits from context-aware memory reading and forgetting, implicitly increasing the model's capacity de- spite a small actual state size, which improves training and inference efficiency. Through extensive experiments, we demonstrate the advantages of GSA in-context recall- inteesiie tasks and in "finetuning pretrained Transformers to RNNs" 4] scenarios.
10

---

Acknowledgments
We would like thank Zhen Qin and Yikang Shen for their insightful discussions, Houquan Zhou and Kazuki Irie for providing valuable feedback on this manuscript.
We gratefully acknowledge the support by LuxiTech for computational resources; and the support by the National Natural Science Foundation of China (No. 62076173, 62476187), the High-level Entrepreneurship and Innovation Plan of Jiangsu Province (No. JSSCRC2021524), and the Project Funded by the Priority Academic Program Development of Jiangsu Higher Education Institutions. Yu Zhang was partially supported by Tencent AI Lab under Wei Bi's mentorship. Songlin Yang was supported by Xianhong Wu Fellowship from MIT.
References
[1] Akyürek, B. Wang, Y. Kim, and Andreas. In-context language learning: Architectures and algorithms. In Proceedings of ICML, 2024. URL http::://pperreiiiieett///ooooo 3S9L
[2] S. Arora, S. Eyuboglu, A. Timalsina, Johnson, M. Poli, J. Zou, A. Rudra, and C. Ré. Zoology: Measuring and improving recall in efficient language models, 2023.
[3] S. Arora, S. Eyuboglu, M. Zhang, A. Timalsina, S. Alberti, D. Zinsley, J. Zou, A. Rudra, and Ré. Simple linear attention language models balance the recall-throughput tradeoff, 2024.
[4] S. Arora, A. Timalsina, A. Singhal, B. Spector. S. Eyuboglu, X. Zhao, A. Rao, A. Rudra, and
C. Ré. Just read twice: closing the recall gap for recurrent language models, 2024. URL
httpps://arxvv. www.eiiii.r.oaaaeeeeeeeeeeoeeeee orgg/b//24////// 05483
[5] J. Ba, G. Hinton, V. Mnih, J. Z. Leibo, and C. Ionescu. Using fast weights to attend to the recent past, 2016.
[6] M. Beck, K. Pöppel, M. Spanring, A. Auer, O. Prudnikova, M. Kopp, G. Klambauer, J. Brand- stetter, and S. Hochreiter. xlstm: Extended long short-term memory, 2024.
[7] A. Bick, K. Y. Li, E. P. Xing. J. Z. Kolter, and A. Gu. Transformers to ssms: Distilling quadratic knowledge to subquadratic models. 2024. URL semanticcccantnhccooco org/CorpusID: 27 1903923
[8] Y. Bisk, R. Zellers, R. L. Bras, J. Gao, and Y. Choi. Piqa: Reasoning about physical common- sense in natural language. In In Proceedings AAAI, 2020.
[9] G. E. Blelloch. Prefix sums and their applications. Technical report, School of Computer Science, Camegie Mellon University, Pittsburgh, PA, 1993. /ttt::ttts:////t CS cmu. edu/ httbb/papers/Ble933 pdf
[10] H. Chen, Z. Liu, X. Wang, Y. Tian, and Y. Wang. Dijiang: Efficient large language mod- els through compact kernelization. ArXiv, abs/2403.19928. 2024. URL https://api semanticscholar. wrwpaarrrrrg/oooooer.... 268793982
[11] T. Chen, B. Xu, C. Zhang, and C. Guestrin. Training deep nets with sublinear memory cost, 2016.
[12] K. Cho, B. van Merriënboer, C. Gulcehre, D. Bahdanau, F. Bougares, H. Schwenk, and Y. Bengio. Learning phrase representations using RNN incoder-decoder for statistical machine translation. In Proceedings of EMNLP, pages 7724-17344 2014. URL https T/aclantholoog.orrg//1---
[13] S. Choi. Cross-architecturee transfer learning for linear-cost inference transformers, 2024. URL
https://axxiv.org/bbb//444 02684
[14] P. Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv: 1803.0555777 2018.

---

[15] T. Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. In The Twelfth International Conference on Learning Representations, 2024. URL / //openreview net/Forum2iddm2222yyya
[16] T. Dao and A. Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality, CoRR, abx/2405.21060, 2024, doi: 10.48550/ARXIV. 2405.2 200.02006.. URL https: org/ /0. 48550/ w4w..
[17] T. Dao, D. Fu, S. Ermon, A. Rudra, and C. Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. In Advances in NIPS, pages 16344- 16359, 2022. ooppr
ppf
[18] P. Dasigi. K. Lo, Beltagy, A. Cohan, N. Smith, and M. Gardner. A dataset of information- seeking questions and answers anchored in research papers. In Proceedings of NAACL, pages 4599-4610, 2021. URL
2021. www.n-aaalaaaaaaaa 365
[19] Y. N. Dauphin, A. Fan, M. Auli, and D. Grangier. Language modeling with gated convolutional networks. In Proceedings of ICMI., pages 933-941, 2017. URI. ht..tttsss//prrceddnggss mlr
press/v70/dauphi 1a..nna...... html
[20] D. Dua, Y. Wang, P. Dasigi, G. Stanovsky, S. Singh, and M. Gardner. DROP: A reading comprehension benchmark requiring discrete reasoning over paragraphs. In Proceedings of NAACL, pages
2019. URL htttps://aclaathhooogg.. hrtpl/oyy.ooo//1//////4/
[21] L. Gao, J. Tow, B. Abbasi, S. Biderman, S. Black, A. DiPofi, C. Foster, L. Golding, J. Hsu, A. Le Noac'h h, H. Li, K. McDonell, N. Muennighoff, C. Ociepa, J. Phang, L. Reynolds, H. Schoelkopf, A. Skowron, L. Sutawika, E. Tang, A. Thite, B. Wang, K. Wang, and A. Zou.
A framework for few-shot language model evaluation, 2023. URL https://zenodo org/ records/10256336
[22] F. Gers, J. Schmidhuber, and F. Cummins. Learning to forget: continual prediction with lstm. In Proceedings of ICANN, pages 550-855, 1999.
[23] A. Gerstenberger, K. Irie, P. Golik, E. Beck, and H. Ney. Domain robust, fast, and compact neural language models. In Proceedings of ICASSP. Barcelona, Spain, 2020.
[24] M. Geva, R. Schuster, J. Berant, and O. Levy. Transformer feed-forward layers are key-value memories. In Pruceedings of EMNLP, pages 14484-5495, Online and Punta Cana, Dominican Republic, 2021. URL https:://cclanthol/ wrwp:/logyyoogg2222///// .em.pp-mmnlppmmim 446.
[25] D. Goldstein, F. Obeid, E. Alcaide, G. Song, and E. Cheah. Goldfinc:: High performance rwkv/transformrr hybrid with linear pre-fill and extreme kv-cache compression. ArXiv, abs/2407.12077 2024. /tthhtsss://api wewantmmnniii remannicscholar.. or wwtp:ooooo/ooooo///o11 2712444::2222445
[26] A. Graves. Generating sequences with recurrent neural networks, 2014.
[27] A. Graves, G. Wayne, and I. Danihelka. Neural turing machines, 2014.
[28] R. Grazzi, J. N. Siems, S. Schrodi, T. Brox, and F. Hutter. Is mamba capable of in-context learn- ing? ArXiv, abs/2402.03170, 2024. URL URthtttps.//pgg..ssaattcccoooot wetemmanniiccool.oorr////oo
267412719
[29] A. Gu and T. Dao. Mamba: Linear-time sequence modeling with selective state spaces, 2023.
[30] A. Gu, C. Gulcehre, T. Paine, M. Hoffman, and R. Pascanu. Improving the gating mechanism of recurrent neural networks. In H. D. III and A. Singh, editors, Proceedings of ICML, pages 3800-3809. PMLR, 2020. URL
mlr wwwppssvslll 199 9/wg20.........// html
[31] D. Hendrycks, C. Burns, S. Basart, A. Zou, M. Mazeika, D. Song, and J. Steinhardt, Mcasuring massive multitask language understanding, 2021. URL https:://arxivorrg/absss 03300
[32] S. Hochreiter and Schmidhuber. Long short-term memory. Neural Computation, 9()):
1735-1780, 1997.
12

---

[33] W. Hua, Z. Dai, H. Liu, and Q. Le. Transformer quality in linear time. In K. Chaudhuri, S. Jegelka, Song, C. Szepesvari. G. Niu, and S. Sabato, editors, Proceedings of ICML, pages 9099-9117. PMIR, 2022. URL https://proceedingss mlr prt.pere///126/ha222 httpp/dddn......pppeee///////aaa html.
[34] K. Irie, A. Gerstenberger, R. Schlüter, and H. Ney. How much self-attention do we need? Trading attention for feed-forward layers. In Proceedings of ICASSP, Virtual only, May 2020.
[35] K. Irie. I. Schlag, R. Csord' as, and J. Schmidhuber. Going beyond linear transformers with recurrent fast weight programmers. ArXiv, 0bs/2106.06295, 2021. URL URL https://appi semanticscholas www.laa.rrr//oooooooee wwt.:oooooooss::23///....o 235417174
[36] K. Irie, I. Schlag, R. Csord' as, and Schmidhuber. A modern self-referential weight matrix that learns to modify itself. In International Conference on Machine Learning, 2022. URL
https://api. semanticscholas ar org/ hegg/CorpusID... 246823084
[37] S. Jelassi, D. Brandfonbrener, M. Kakade, and E. Malach. Repeat after me: Transformers are better than state space models at copying. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL https: httpp:://ppeneee //tperri/pperreviiee.. net/ mem//foorumm
[38] Jelassi, D. Brandfonbrener, M. Kakade, and E. Malach, Repeat after me: Transformers are better than state space models at copying, 2024.
[39] A. Q. Jiang, A. Sablayrolles, A. Mensch, C. Bamford, D. S. Chaplot, D. de las Casas, F. Bres- sand, G. Lengyel, G. Lample, L. Saulnier, I.. R. Lavaud, M.-A. Lachaux, P. Stock, T. L. Scao, T. Lavril, T. Wang, T. Lacroix, and W. E. Sayed. Mistral 7b, 2023.
[40] M. Joshi, E. Choi, Weld, and L. Zettlemoyer. TriviaQA: large scale distantly supervised challenge dataset for reading comprehension. In R. Barzilay and M.-Y. Kan, editors, Proceed- ings of ACL, pages 1601-1611, Vancouver, Canada, July 2017. doi: 10. 1...8///111117-11477 URL https: https////caaathhoegoooo http:///aaaaooollooo///////---// org/P17-1147
[41] J. Kasai, H. Peng, Y. Zhang, D. Yogatama, G. Ilharco, N. Pappas, Y. Mao, W. Chen, and N. A. Smith. Finetuning pretrained transformers into RNNs. In M.-F. Moens, X. Huang, L. Specia, and S. W.-t. Yih, editors, Proceedings of EMNLP, pages 10630-10643, 2021. URL https://acclanthology grg/ 2021 emnlp-main. 830
[42] A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret. Transformers are RNNs: Fast autore- gressive transformers with linear attention. In H. D. III and Singh, editors, Proceedings of (CMM/,,pages 5156-5165. PMIR, 2020. URL. katharopoulos20a..no
mlr 1 press/v119/
[43] T. Katsch. Gateloop: Fully data-controlled linear recurrence for sequence modeling, 2024.
[44] T. Kočiský, Schwarz, P. Blunsom, C. Dyer, K. M. Hermann, G. Melis, and E, Grefenstette. The NarrativeQA reading comprehension challenge. TACL, pages 317-328, 2018. URL ht https:://aclanthology htt.loggg.orr///// -
[45] Krotov and J. Hopfield. Large Associative Memory Problem Neurobiology and Machine Learning, 2021. URL
06996
[46] T. Kwiatkowski, J. Palomaki, O. Redfield, M. Collins, A. Parikh, C. Alberti, D. Epstein, I. Polosukhin, Devlin, K. Lee, K. Toutanova, L. Jones, M. Kelcey, M.-W. Chang, A. M. Dai, J. Uszkoreit, Q. Le, and S. Petrov. Natural questions: A benchmark for question answering research. TACL, pages 452-466, 2019. URL httpss///allanhhology wttpa/ooyyyoo/119111999900
[47] O. Lieber, B. Lenz, H. Bata, G. Cohen, J. Osin, I. Dalmedigos, E. Safahi, S. Meirom, Y. Belinkov, Shalev-Shwartz, O. Abend, R. Alon, T. Asida, A. Bergman, R. Gloz- man, M. Gokhman, A. Manevich, N. Ratner, N. Rozen, E. Shwartz, M. Zusman, and Y. Shoham. Jamba: A hybrid transformer-mamba language model, 2024. URL, https
19887
13

---

[48] B. Liu, R. Wang, L. Wu, Y. Feng, P. Stone, and Q. Liu. Longhom: State space models are amor- tized online learners. ArXiv, abs/2407.14207, 2024. pttttps:s/api.ssmattissssott
org/CorpusTD: 271310065 27111165
[49] C. Lockard, P. Shiralkar, and X. L. Dong. OpenCeres: When Open Information Extraction Meets the Strut-rtre Web. In Burstein, C. Doran, and Solorio, editors, Proceedings of NAACL. pages 3047-3056, Minneapolis. Minnesota, 2019. doi: 10.18653/v1/N19-111 htL.ltttps//allaatilloo..oo wtwp://ooooooor//-1113110
[50] I. Loshchilov and Hutter. Decoupled weight decay regularization, 2019.
[51] X. Ma, X. Kong. S. Wang, C. Zhou, May, H. Ma, and L. Zettlemoer. Luna: Lin- ear unified nested attention. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in NIPS, volume 34, pages 2441-2453, Curran Associates. Inc., 2021. URL https::/ppooceediinng neur cntuurpps.cooo http///aerr.llle///ppr/////1///1 1441900666111100000022244411f113
[52] H. H. Mao. Fine-tuning pre-trained transformers into decaying fast weights. In Y. Goldberg. Z. Kozareva, and Y. Zhang, editors, Proceedings of EMNLP, pages 10236-10242, Abu Dhabi, United Arab Emirates, 2022. UtL//tttss///aclaannooooyyy org. emnle-main...n... 697
[53] E. Martin and C. Cundy. Parallelizing linear recurrent neural nets over sequence length. In Proceedings of ICLR, 2018. URL htt/htttss//oppeeeiieee
[54] J. Mercat, I. Vasiljevic, S. Kch, K. Arora, A. Dave, A. Gaidon, and T. Kollar. Linearizing large language models, 2024.
[55] S. Merity, C. Xiong. J. Bradbury, and R. Socher. Pointer sentinel mixture models, 2016.
[56] T. Munkhdalai, A. Sordoni, T. Wang, and A. Trischler. Metalearned neural memory. ArXiv, abs/1907.09720. 2019. URL httth:://///pp semanicsccoolar org/CorpusID: 1981 79407
[57] M. Oren, M. Hassid, Y. Adi, and R. Schwartz. Transformers are multi-state rnns, 2024.
[58] R. Pang, A. Parrish, N. Joshi, N. Nangia, Phang, A. Chen, V. Padmakumar, J. Ma, J. Thompson, H. He, and S. R. Bowman. Quality: Question answering with long input texts, yes!, 2022. htRLhhtps::// //txxv/aaivvvv//a http///axxv.orggabb///// wwtp/b///////////o/// 08608
[59] D. Paperno, G. Kruszewski, Lazaridou, N. Pham, R. Bernardi, S. Pezzelle, M. Baroni, G. Boleda, and R. Fernández. The lambada dataset, 2016.
[60] G. Penedo, Q. Malartic, D. Hesslow, R. Cojocaru, A. Cappelli, H. Alobeidli, B. Pannier, E. Almazrouei, and J. Launay. The RefinedWeb dataset for Falcon LLM: outperforming curated corpora with web data, and web data only. arXiv preprint arXiv:2306.00 16, 2023. URL WtR/hhtps:///rxiiv http://wwwrr//bb///////a//
01116
[61] B. Peng. D. Goldstein, Q. Anthony, A. Albalak, E. Alcaide, S. Biderman, E. Cheah, X. Du, T. Ferdinan, H. Hou, P. Kazienko, K. K. GV, J. Kocoń, B. Koptyra, S. Krishna, R. M. J. au2, N. Muennighoff, F. Obeid, A. Saito, G. Song, H. Tu, S. Wožniak, R. Zhang, B. Zhao, Q. 7hao, P. Zhou, J. Zhu, and R.-J. Zhu. Eagle and finch: Rwky with matrix-valued states and dynamic recurrence, 2024.
[62] H. Peng, N. Pappas, D. Yogatama, R. Schwartz, N. Smith, and L. Kong. Random fea- ture attention. In Proceedings of ICLR, 2021. URL htt.ttts://oppeeeeiewww httt::://ppprreeeeeeee////rrrrr QtTKTdVrFBB
[63] H. Peng, J. Kasai, N. Pappas, D. Yogatama, Z. Wu, L. Kong, R. Schwartz, and N. A. Smith. ABC: Attention with bounded memory control. Proceedings of ACL, pages 7469-7483, 2022. httphttts///aaanoolloogg// wwwp:/loooooo.gg/ee0 wwt.:org./2..... acl -log 515
[64] R. Pope. S. Douglas, A. Chowdhery, J. Devlin, J. Bradbury, A. Levskaya, J. Heek, K. Xiao,
S. Agrawal, and Dean. Efficiently scaling transformer inference, 2022.
[65] S. Pramanik, Flelimy, M. C. Machado, and A. White. Recurrent linear transformers, 2023.
URL hww.h:tts//eeee //arxiv..iivvvv//a org 2b10/ 15719
15719
14

---

[66] Z. Qin, X. Han, W. Sun, D. Li, L. Kong, N. Barnes, and Y. Zhong. The devil in linear transformer. Y. Goldberg, Z. Kozareva, and Y. Zhang, editors, Proceedings of EMNLP. pages 7025-7041, Abu Dhabi, United Arab Emirates, 2022. doi: 10. 10116635/1/202222mmmmmmm...... URL http:/tttss//claaiiioooooooooo http://ooo.....o/////.mmmm..aaaa emnlp-main,
[67] Z. Qin, S. Yang, and Y. Zhong. Hierarchically gated recurrent neural network for sequence mod- eling. In Advances in NIPS, 2023.
[68] Z. Qin, D. Li. W. Sun, W. Sun, X. Shen, X. Han, Y. Wei, B. Lv, X. Luo, Y. Qiao, and Y. Zhong. Transnormerllm: A faster and better large language model with improved transnormer, 2024.
[69] Qin, Yang, W. Sun, X. Shen, D. Li, W. Sun, and Y. Zhong. Hgmn2: Gated linear rnns with state expansion, 2024.
[70] P. Rajpurkar, Jia, and Liang. Know What You Don't Know: Unanswerable Questions for SQuAD. In Proceedings of ACL, Melbourne, Australia, 2018. Association for Computational Linguistics.
[71] L. Ren, Y. Liu, Y. Lu, Y. Shen, C. Liang, and W. Chen. Samba: Simple hybrid state space models for efficient unlimited context language modeling. CoRR, abs/2406.07522,, 2024. doi: 10.48550/ARXXV12406..7777
URL https: /tttpp://doi.org 01011...8////rii 2406. 07522
[72] A. Rush. Torch-struct: Deep structured prediction library. In A. Celikyilmaz and T.-H. Wen, editors, Proceedings of ACL, pages
Online, July 2020. doi: 10.18653/v1/2020.
acl-demos.38 URL https: http:/ss/aaacaaatoooog...r//////
-emos.. 38
[73] Sakaguchi, R. Bras, C. Bhagavatula, and Y. Choi. Winogrande: An adversarial winograd schema challenge at scale, 2019, httphttps//raaioorg//////
10641
[74] I. Schlag and J. Schmidhuber. Gated fast weights for on-the-fly neural program genera- tion. In Proceedings of ICLR, 2017. URL /tt.///tsss//ppppps.ttttttssoooa www.:/sssataaaiiaaaa.a.ooo wwtp:/oorrrroo/ooooooooooo 216094255
[75] I. Schlag, K. Irie, and J. Schmidhuber. Linear transformers are secretly fast weight program- mers. In M. Meila and T. Zhang, editors, Proceedings of ICML, pages 18-24 Jul 2021. URL
PMLR,
[76] I. Schlag, T. Munkhdalai, and Schmihhuber. Learmi Learning associative inference using fast weight memory, 2021. URL /RRL/ttpss//arriv..rr///////////
07831
[77] J. Schmidhuber. Learning to control memorie::s An alternative to dynamic recurrent networks. Neural Computation, 4(1):131-139, 1992.
memories:
[78] N. Shazeer. Glu variants improve transformer, 2020.
[79] D. Soboleva, F. Al-Khateeh, R. Myers, J. R. Steeves, J. Hestness, and N. Dey. Slimpa- jama: A 627b token cleaned and deduplicated version of redpajama, 2023. URL https / //huggingf ace. co/
[80] J. Su. Y. Lu, S. Pan, A. Murtadha, Wen, and Y. Liu. Roformer: Enhanced transformer with rotary position embedding, 2023.
[81] S. Sukhbaatar, a. szlam, J. Weston, and R. Fergus. End-to-end memory networks. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett, editors, Advances in NIPS. Curran Associates, Inc., 2015. URL
[82] Y. Sun, L. Dong, S. Huang. S. Ma. Y. Xia. J. Xue, J. Wang, and F. Wei. Retentive network: A successor to transformer for large language models, 2023.
[83] Y. Sun, L. Dong, Y. Zhu, S. Huang, W. Wang, S. Ma, Q. Zhang, J. Wang, and F. Wei. You only cache once: decoder architectures for language models, 2024. URL
httpp:://axxiv.org/bb/////// 05254
15

---

[84] Y. Sun, X. Li, K. Dalal, J. Xu, A. Vikram, G. Zhang, Y. Dubois, X. Chen, X. Wang, O. Koyejo. T. Hashimoto, and C. Guestrin. Learning to (leam test time): Runs with expressive hidden states. ArXiv, 0bb/2200.004620 2024. htt.h:ttps:///pt www.nti wtg/Corppsiii 271039606
[85] M. Suzgun, N. Scales, N. Schärli, S. Gehrmann, Y. Tay, H. W. Chung, A. Chowdhery, Q. Lc, E. Chi, D. Zhou, and J. Wei. Challenging BIG-bench tasks and whether chain-of-thought can solve them. In Findings of the ACL, pages 13003-13051, Toronto, Canada, 2023. URL https://aclanttollee .og.loloooooooooooooooooo org 2023. findings acl.824.
[86] H. Touvron, T. Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, A. Rodriguez, A. Joulin, E. Grave, and G. Lample. Llama: Open and efficient foundation language models, 2023.
[87] H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, D. Bikel, L. Blecher, C. C. Ferrer, M. Chen, G. Cucurull, D. Esiobu, J. Fernandes, J. Fu, W. Fu, B. Fuller, C. Gao. V. Goswami, N. Goyal, A. Hartshorn, S. Hosseini, Hou, H. Inan, M. Kardas, V. Kerkez, M. Khabsa, I. Kloumann, A. Korenev,
Koura, M.-A. Lachaux, T. Lavril, J. Lee, D. Liskovich, Y. Lu, Y. Mao, X. Martinet, T. Mihaylov, P. Mishra, I. Molybog,
Nie, A. Poulton, Reizenstein, R. Rungta, K. Saladi, A. Schelten, R. Silva, E. M. Smith, R. Subramanian, X. E. Tan, B. Tang, R. Taylor, A. Williams, J. X. Kuan, P. Xu, Z. Yan, Zarov, Y. Zhang, A. Fan, M. Kambadur, S. Narang, A. Rodriguez, R. Stojnic, S. Edunov, and T. Scialom. Llama 2: Open foundation and finc-tuned chat models, 2023. URL https:
argg/bbs/2300 09288
[88]
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in NIPS. Curran Associates, Inc., 2017. URL https::/ neur ips. httppccpcppp....ppppppppppppp..p wttpc/aaee..llle//pppeeeeeeeeeee 3f5ee243547dees
Paper.
[89] R. Waleffe, W. Byeon, D. Riach, B. Norick, V. Korthikanti, T. Dao, A. Gu, A. Hatamizadeh, Singh, D. Narayanan, G. Kulshreshtha, Singh, J. Casper, J. Kautz, M. Shoeybi, and B. Catanzaro. An empirical study of mamba-based language models, 2024. URL https
//arxiv.org/abbeitb
07887
[90] J. Wang, D. Paliotta, A. May, A. M. Rush, and T. Dao. The mamba in the Ilama: Distilling and accelerating hybrid models, 2024. URL httts:://axxii http:://axxvvvorgr./aa www.org///bbb.rge.eeee 15237
[91] K. Wen, X. Dang, and K. Lyu. Rnns in-context retrieval. ArXiv, abs/2402.18510, 2024. URL htttps://app/ semantisccolla htg/CorppusD:22884411
transformers (yet): The key bottleneck
[92] B. Widrow and M. E. Hoff. Adaptive switching circuits. 1988. semanticscholar- arg/CorpusID: 60830585
URL https://api
[93] E. Wu, K. Wu, R. Daneshjou, D. Ouyang, D. E. Ho, and J. Zou. How medical AI devices are evaluated: limitations and recommendations from an analysis of FDA approvals. Nature Medicine, pages 582-584, 2021. URL httpt:t//t...oorrrrrrr
1038/ 01588/441511-222-11112-22-
[94] W. Xiong, J. Liu, I. Molybog, H. Zhang, P. Bhargava, R. Hou, L. Martin, R. Rungta, K. A. Sankararaman, B. Oguz, M. Khabsa, H. Fang, Y. Mehdad, S. Narang, K. Malik, A. Fan, S. Bhosale, S. Edunov, M. Lewis, S. Wang, and H. Ma. Effective long-context scaling of foundation models, 2023. URL/htts:://ara/
org/ abs /339 16039
[95] S. Yang and Y. Zhang. FLA: A Triton-Based Library for Hardware-Efficien Implementa- tions of Linear Attention Mechanism, 2024. URL /tttps://githhbb...ttt com/ sussssssstcsoniii flash-linear- attention
[96] S. Yang, B. Wang, Shen, R. Panda, and Y. Kim. Gated linear attention transformers with hardware-efffcient training. In Proceedings of ICML. PMLR, 2024.
16

---

[97] S. Yang, B. Wang. Y. Zhang, Y. Shen, and Y. Kim. Parallelizing linear transformers with the delta rule over sequence length. CoRR, 0ed/2406.066844, 2024. doi: 10.48550/ARXIV:24006 06484. https://:/// org/10. 48550/arXiv 2406.00220000 06484
[98] S. Yang. B. Wang, Y. Zhang, Y. Shen, and Y. Kim. Parallelizing linear transformers with the delta rule over sequence length. ArXiv, abs/2406.064844, 2024. URL https://api
semanticscholarr. .rg. Ttrg/CorpusID:::0 270371554
[99] R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi. Hellaswag: Can machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019.
[100] J. Zhang, Jiang. J. Feng, L. Zheng, and L. Kong. Cab: Comprehensive attention benchmarking on long sequence modeling. ArXiv, abs/2210.07661, 2022. URL https
//api. www.:esemaiiscchol..o. wwwpaanttccccola.oooo/o/ooo///// 252907545.
[101] M. Zhang, K. Bhatia, H. Kumbong, and Ré. The hedgehog & the porcupine: Expressive linear attentions with softmax mimicry, 2024.
[102] P. Zhang. G. Zeng, T. Wang, and W. Lu. Tinyllama: An spence small language model, 2024.
[103] W. Zhang and B. Zhou. Learning to update auto-associative memory in recurrent neural networks for improving sequence memorization. ArXiv, nos//779.06493, 2017. URL https:
//api. www.sssematiicccoo...o.o hwwpaaaiiiiiccaaaaoooo//orrrssps wotp:ooooooooo 22458497
[104] Y. Zhang and D. Cai. Linearizing transformer with key-value memory. In Y. Gold- berg, Z. Kozareva, and Zhang, editors, Proceedings EMNLP, pages 346-359 Abu Dhabi, United Arab Emirates, 2022. doi: 10. 18653/11/202..eel 16653//1102222emmpmmmm....22 URL https //aclanthology' wwwpp//........ooooo..oooo.o.... wttg/2022.emmlp-mme. emt.:2mammllllm.i.... 24
[105] M. Zhong, Yin, T. Yu, A. Zaidi, M. Mutuma, R. Jha, H. Awadallah, A. Celikyilmaz,
Y. Liu, X. Qiu, and D. Radev. QMSum: A new benchmark for query- multi-domen meeting summarization. In Proceedings of NAACL, pages 5905 5921, Online, 2021. doi: 10. 12228865/11/2221mmm w.waa20-..aaammmmmmmmm maa.c-mmin..7777 URL htLLtttss///ccattttttt
org/ 2021 www.aaalllma. mraac-mmain
472
[106] G.-B. Zhou, J. Wu, C.-L. Zhang, and Z.-H. Zhou. Minimal gated unit for recurrent neural networks, 2016.
17

---

A Linear Attention and its Chunkwise Form
alternative to resolve the quadratic complexity of self-attention (SA). The key idea is to use the kernel trick, which replaces softmax with decomposable kernel function, resulting the following parallel
Linear Attention (LA)
42 66 68
emerges
O (()(Q)(())) M)V.
(10)
where Φ IRd Tn functions as feature mapping applied to each input. Unfolding Eq. 10 have W De, Wa W, E ER",
42 kit,
Σ (k)(())) o(g,)= S Σ o(k) @ ==1
(11)
= v fkk 9.)
o(q.).
1=1
means outer product operation. is clear that by leveraging the associativity, LA admits simple recurrent updating rules with matrix-valued hidden states S. E mmd
O S ((q:); S, S + <(k,) vr.
(12)
By reserving bounded m memory slots only, the overall computation complexity is reduced from d) to (T'md) When the sequence length the complexity, and LA can be much more efficient than its counterpart with quadratic complexity.
>> m, d, the md factor has minor impact on
During inference, LA enjoys the merits of RNNs, which only need to maintain ((md) hidden memories, helping avoid the memory-cost KV cache management
SA mechanisms. However, Eq. [2 employs simple additive updating rule and can be hard to "forger" unrelated information if necessary [62], making the limited memory states vulnerable to be chaotic.
Gating mechanism has played a key role in classical RNNs 32 mechanism to control the information flows in the network and help read and write from the memory selectively. [82] propose to apply a data-independent gate to LA, significantly narrowing the gap
12 which serves as
between LA and SA: S, AS + ((k,) @ vt. À [0. 1] is a non-learnable scalar. Recent
work [96 further imposes finer-grained data-dependent gate:
(13) where each a E 10 from A == {ax {a E [0. 1| " is dependent on the input. Alternatively, we can couple the key values with the forget gates by allowing o(k,) a in spirit of [[2 [06 and 669 which reduces the number of parameters and improves efficiency accordingly.
S, Diag S, + ó(k,) v.,
A.1 Hardware-Efficies Training
Despite the theoretical advantages of linear complexity, the recurrent form of Eq. 12 is still inefficient during training. Such recurrent computation prevents the full utilization of modern GPU parallelism over sequence lengths [53] 72). On the other hand, the parallel form (Eq, [T) can be parallelized in similar vein as flash attention [7 However, due to the existence of the casual mask M, we can not rearrange its computation order by KV first, so that the parallel form still adheres to the quadratic complexity, which can hardly be scaled very-long training context (e.g., sequences with more than 8K tokens).
Chunkwise form recurrences have been carried forward by [2). and achieve good trade-off between the recurrent and parallel forms. (96) further disclose that the glement-wiss gating of Eq. 13 also satisfies the associative property required by parallel scan [9] and derive a parallelized chunkwise gated linear attention in similar vein. The key idea is to partition the sequence into N = [ chunks
of size C' with Q = 4ic, 9ic+1..... 9ic-c. and so forth for K(t), V E Bxxd An E RC Firstly, unrolling the i-th hidden state the t-th chunk in Eq. [3]
we get
S pp Diag (A1))) S1g,611 + & (Kigja) & Vig
(14)
Diag
II
A M)
S S + Σ
(K (.a ④
II Aq
k=
j=
Whhere is normalization term in vanilla LA similar to softmax, potential gradient explosions.
666 reveal that removing it could avoid
18

---

We write the last hidden in the chunk Stgcc Stg) interchangeably for simplicity. Define A
pls E [0, 1] as the cumulative decay from the start of chunk to i, and likewise A Ap), = = Ans E [o, 1j from i + to the end of the chunk, then
S ( Diag [1..0) (K1) A14) V(4
(15)
A, A can be absorbed into Q K first: [d (QQ() A [t)- K ((K (1) (Atg/
A
m.)))
Combining them with Eq. 10 and Eq. 14 we derive the following vectorized updating rules
Q (c) S [t
(g о MMa)
V(4
(16)
The first term is referred to the inter chunk part and the second term is the intra chunk part. The process get this intra part little more involved the cumulative productions of A Atg. is greater than 1, which can lead to numerical instability. [96] deal with this issue by proposing secondary-chunking strategy, and we refer readers to their paper for more details.
Hardward considerations Modern GPU architectures, such as the NVIDIA A100, offer highly optimized matrix multiplication (matmul) operations through specialized Tensor Cores, achieving up to 16x higher throughput than non-matmul operations [17]. However, this incurs IO overheads due data transfer from slower, off-chip global high bandwidth memory (HBM) on-chip shared memory (SRAM). The chunkwise form balances I/O and computation complexity tradcoffs. As shown in Eq 16. improves parallelism over the sequence dimension while reducing non-matmul FLOPs greatly. Also, the chunk recurrent updating conducts the query and hidden states reduction in an online manner, requiring only O(Ndm) hidden states materialized into HBMs, so that it can significantly reduce the memory/10 overheads. While LA enjoys much lower overall running FLOPs than SA, the chunkwise form displays practical significant wall-clock speedup against SA, due to its Bardware-efficien implementations [95].
B Details for GSA
Opi
[2]
Beyond the recurrent GSA form pro- vided in Figure. [ we give de- tailed, hardware-efficient procedures for
%%
%110
%
the forward and backward passes of Gated Slot Attention (GSA) in Algo- rithm [ For simplicity, we define
o
Fpp 1/3
A az [ =00, 10.m and
{1 a- x] The algo- rithm demonstrates that GSA can be modeled two-pass GLA, illus- trated in Fig.
Sfa)
Qp
K K
Kiyji Ara
In the preprocessing step. pre- compute the chunkwise cumulative sum of the forget gate, resulting in A. Subse- quently, A along with the queries, keys, and values are passed to engage in two
Figure 5: Diagrams of the recurrence and updating rules in Gated Slot Attention. The outputs of the first pass taken as queries the second pass. 8 query nodes
key/value nodes
output nodes
recurrent hidden states
GLA passes. For each chunk of size C, we define A A as in Eq. and Eq. 16
In the first pass, A, 'A is absorbed into K K
Q Q 3
A f K
K [)) o (0) Ag..).
then Q and K function as usual queries and keys, and the slot representations I serve the value vectors.
Q
S 11-11
Q is]
K kT)
M)
RRCm
[
Outer
We use different notations from those presented in oo enhance clarity in the chunkwise updating rules. The output 0 is decomposed into the inter-chunk recurrence and intra-chunk parallel computations (96).
19

---

Algorithm 1 Hardware-Effhcient Gated Slot Attention
FORWARDPASSII,
BACKWARDPASSIQ A. O 20) Divide Q. K, 00 R 1 E.A.EEE Tm
Divide Q. K. PEEEEXX 2 1, x
,,A,
into N = =1 blocks
C chunk size
into N = =
blocks
> C is chunk size
function chunk ())
function gea, bea_b K, V. dO, GATE,
parfor 1.
[n] t0 SRAM
Onehip do
M RCXC
n Store d$ in HBM as (S
reverse
Store [m <- compus (A A HBM
Ininalize ds 0
return + A 10)....
[
Load (1) A (c:) (x) SRAM
On chip:
function fodlQ,
A, BATR On chip: construct cuusal mask M xexe N.NNA
IFGATE_xthae then
Diag( 7 [1), + (9p1) @ A (m0ml
ia HAH [nm]
Initialize
dSDiag(( + (d0mm @ A(a))
Load K/nl. V(n: A1) A SRAM
parfor n +- 1, do
On chip: [1)] A /n).0/ (n)
Load Qjuj. VIn xxxd
if HATAR x then
([n].d[[a] gdxa to SKAM Oa chip:
Diag (0) + (14p)) Amm) ((n)
Recompune A q.1)) R/nl. [nn..
ehs
8 SDingl m),0)) Kin] (V) (n] Ann))
if (FGATE_ x then
dP - (d0 V (n) @
parfor n - 1.
aQ - (d0p] + dPk (1) @ Apm
[nn-- Kin). V/nl Spul. Toy Ain SRAM
dK (Vinjds + aP Qm1) @ A|n
On chip.
K(m:dSpu] pm
((m) + [n] [(n)
de
[n]] [(n] (AA [n] A (m)
dP (d0 V
[n] - Q1,]
== q/nl RT) M) Vial
aQ d0/m| dPK
dK 1/ + aap qu1
ebe
[m] + Vin 1m (mjc)
av + dS P p Tno
[m] (n- Qid K/ e M) Vnn
Wine uQQ.Q,KK.. w dQ (m) [V
[a) 011 [m)
return dQ pj.......
NI.
dV WV |.....
Store O [p])
Recompute 7. s" dQ dr" dV + g b V. d, Faice) d +- dsoftmax(O" aQ") aQ dK dI" bna K. ss 20 True) al -- ar* +1p" AA reversed 0 pQ K aK+
return Ny. St.........
I c- chank AA)sun(A)
s-ootmaa gradients
pregncessing
0 st e- Qudi K, 1, Falsw)
q" - =oftmax(O**
0, gaat
True)
0 0010
vV)
return
return JQ, dV,
In the second pass, the output O* from the first pass, after the application of the softmax function, serves as the queries Q",
Q"
softmax O²
and I/V are used as the key/value vectors, respectively. The final GSA output O is obtained as follows:
S [-11 ((Q), M)MEE R
[
Unlike in the first pass, A. AA is absorbed into V. rather than Q.K.
During the backward pass, computing the gradients of Q, K, V, A involves variables already com- puted in the forward pass. However, directly saving all intermediate results can pose severe challenges for memory management. To address this issue, we adopt gradient checkpointing 11 to trade off memory consumption for recomputation. In addition to the input Q K, V, V, A, we selectively save only the output of the first GLA pass, which significantly reduces memory consumption (Figure 4b).
Similar to the forward pass, the backward pass involves two GLA backward passes well, but in the reverse order. The final gradient dI is obtained by combining the gradients from these computations, ie., dI = a + dI". The forget gate gradient can be decomposed into two parts: Q O dQ - K o dK and O @ dO V ② dV (cf. $C in (96). The reversed cumulative sum in the backward pass corresponds to the cumulative
computed in the preprocessing step of the forward pass.
We provide a Py Torch implementation for the above algorithm with chunkwise parallelism in Listing
20

---

def ksd_fwd k, v, E, C):
q/k/v:
query, key, value of shape INC, C, K
local cumulative product of forget gate log space
C:
chunk size
# NC: number of chunks
11
# query/key head dimension
12
# V: value head dimension
13
NC, C, K, V shape, v. shape[111
14
# [K, V]
15
s q new_zeros(K, V)
16
# [NC, C, V]
17
o torch. eeaphheettyyeleeeeee
18
19
for i in range(0, NC):
28
[C, KIVJ chunking
21
c_9, c_k, c_v, c_g q[i], k[i], v[iJ, g[i] # the last of each chunk
22
23
c_gn = c_g[-1]
inter-chunk w/ matmul
25
c_vg, c_gn CLV (c_gn c_g) exp(), e_gn.exp()
26
# [C, V]
27
c_o_inter (c_q s) e_g.exp()
28
# hidden state update c_gn t.k.(() c_vg
29
38
31
# intra-
32
# [C, C]
C_A c_q c_k.t()
34
[c, V]
35
c_o_intra torch. zeros zeros_leleeeeee _llkkec_kv)
36
for in range(0, C // 16):
slice(j * 16, [16, KIV] subchunking
16 16)
38
39
s_A, s_v, s_g c_A[t], c_v[t], c_g[t] s_o gerew_ V)
48
41
42
inter-subchunk w/ matmul s_gn s_g[ø]
43
44
for si range(0, j):
45
u= ssice(si * 16, si * 16 16)
46
$_0 += s_A[;, u] @ (c_v[u] (s_gn c_g(ul) exp())
47
$_0 (s_g s_gn) exp((
48
intra-subchunk w/o matmul for si in range(16):
49
58
for sj in range(si + 1):
s_o[si] += s_A[si, 16 sj] s_v[sj] * (s_g[si] s_g[sj]).exp) c_o_intra[t] s_o
52
53
# [C, V]
54
o[i] = c_o_inter c_o_intra
55
return
57
58 def gsa_fwd. fwd_v(q, k, v, g, C):
59
NC, C, K, V shape, v. 1h (K,,,, v)
61
torch. wow.h.empty_liiii. empty_like(v)
62
63
for in range(0, NC):
21

---

64
# [C, KJV] chunking
c_q, c_k, c_v, c_g q[i], k[il, v[il, g[i] the last g of each chunk
66
67
c_gn = c_8[-1]
68
# inter-chunk w/ matmul
c_9g, c_kg, c_gn c_q * c_g.exp(), c_k (c_gn exp).exp() exp [C, v]
71
c_n_inter c_ag e
72
# hidden state update
73
c_gnl: None] * s c_kg. c_v
75
intra-chunk
76
c_A q. n.nne_zeros(( C)
77
for in range(0, C // 16):
18
slice(j * 16, * 16 + 16)
[16, KIV] subchunking
88
s_q, s_k, s_g c_q[t], c_k[t], c_g[t] $_A q. ger 16)
81
82
intra-subchunk w/o matmul for si range(16): for sj in range(si + 1):
84
85
86
s_A[si, sj] torch. sun(s_q[si] s_k[sj] * (s_g[si] s_g[si]) exp()) c_A[t, t] = S_A
87
inter-subchunk w/ matmul s_gn s_g[e]
89
98
s_9g = s_q (s_g
91
for si in range(0, i):
slice(si * 16, si 16 16)
93
c_Ast, u] s_ag @ (c_k[u] (s_gn c_g(u]).expp to
94
c_o_intra c_A c.v
95
# [c, V]
96
o[i] = c_o_inter c_o_intra
return
98
99
168
def gsa(q, k, v, s, g):
M = shape
# reshape each input [NC, KIV]
9, k, v, s, g = (1ambaa x: x. view(-1, c, shappe[[1,)), (q, k, v, s, 8))
164
# local compute of cumulative product of decay
# [NC, C, K]
g g. c1)
gsa_ fwq_ffd_k(q, s, g, M)
188
qv ok. soft
109
o psa_fwd_v(qy, v, 8, M)
return view(T, -1)
Listing 1: Pseudo PyTorch-style code snippet for GSA with chunkwise parallelism. For brevity, omit the dimensions of batch size and number of heads. Notably, unlike Algorithm we obtain the intra outputs via secondary chunking strategy in Line 31-52 and Line 75-94, as utilized by GLA 196, to ensure numerical stability.
C Experimental Setup
C.1 Language Modeling
We compare GSA with the following strong Transformers with modern architectural recipes as well as other recent subquadratic architectures:
Llama-like architectures that enhance the vanilla Transformer by using Rotary position embeddings (80) and GLU (78).
Xfmr++ 86]:
22

---

Mamba
29
State-space models with data dependennnn decay.
RetNet 18
Linear attention with non-learnable, dula-independenl head-wise decay and rotary
embedding.
GLA |6 Linear attention with elementwise data-dependent decay.
HGRN2
69
Gated Linear RNN with state expansion, or GLA with improved parameterization.
Setup. For fair comparison, all models are trained from scratch with the same training recipes. We utilize a subset of 100B tokens picked from the Slimpajama dataset [79]. The input tokens are processed using the Mistral tokenizer 39 We use Adam W (50) with a weight decay 0.01 as the optimizer. During training, the learning rate is first warmed up to 3 x 10-4 in the first 1B tokens, and then decayed to 3 10-8 gradually with a cosine schedule. The number of attention heads set to and 5 for 1.3B and 2.7B models, respectively. The number of memory slots is uniformly set to 64 for all models. We utilize the open-sourced Triton-based library FLA [ to run all compared models.
We ran all models on 32 Nvidia H800 GPUs. To facilitate distributed training and accelerate the process, we utilized the DeepSpeed framework and fused all necessary modules, including ROPE, cross-entropy, and LayerNorm, following the practice of [102). The training of GSA model with 2.7B parameters took approximately 2 days, while the 1.3B model required day to complete training.
Remark on state size. Let the model dimension be denoted as d. Mamba expands the value projection to 2d and uses a state expansion ratio of 16, resulting in a state size of 32d per layer. Since Mamba also replaces the FFN with a Mamba layer, this effectively doubles both the number of recurrent layers and the state size, leading total recurrent state size of 64Ld.
Similarly, RetNet expands the value projection to 2d and sets the head dimension of queries/keys to be half that of the value head dimension. RetNet also reduces the number of heads to increase the head dimensions of queries and keys. We fix the query/key head dimension to 256 and adjust the number of heads accordingly, resulting in a recurrent state size of 512d per layer and 512Ld in total.
GLA does not expand the value projection but reduces the head dimensions of queries and keys to half of the value head dimension to save parameters for the Swish output gate, ensuring each layer contains 4d parameters. We fix the query/key head dimension to 256 and adjust the number of heads accordingly. resulting in recurrent state size of 256d per layer and 256Ld in total.
HGRN2 follows a similar approach GLA but without the Swish output gate, keeping the head dimensions of queries/keys and values equal, as in standard softmax attention, while still retaining 4d total parameters per recurrent layer. We set the head dimension to 128, resulting in a recurrent state size of 128d per layer and 128Ld in total.
GSA maintains hidden states for both keys and values, so each layer contains a recurrent state size of 2 x 64 x d. We fix the state expansion (i.e., number of slots) to 64 12 resulting in a total recurrent state size of 128Ld.
"/nttps://huggingffce http://nnaaaa..ooo/sssrraaaaa/aa co/misstralai//isttall------
Note that in this case, the number of heads is independent of the state expansion ratio
23

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: This paper's contributions and scope are reflected in abstract and introduction part clearly.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations of this work in Section 6.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]

---

Justification:
This paper does not include theoretical results that require full proof.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
The paper provides sufficient details on hyperparameters and training proce- dures to reproduce the results supporting its main conclusions.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
25

---

Answer: [Yes]
Justification: The code, data and pretrained models are publicly available at GitHub and Huggingface. The training processes are reproducible following the guidance in the paper.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co cc/
oublic//uiies/Cooeeuummi Suw.iiiiioiooooieyy for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run reproduce the results, See the NeurlPS code and data submission guidelines (https
sstoomisssinPoo ssionPolicyyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We have detailed all the training and evaluation settings before the main results in the experimental part.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: We do not have enough resources to obtain error bars as running the experi- ments multiple times is computationally expensive due to the large model size.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
26

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide information of GPU type and number of GPUs used for running our experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https:////rpiips.... httph//rrppp..pppiii/c/ttcsseeee
Answer: [Yes]
Justification: This work follows the NeurIPS Code of Ethies.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
We foresee no potential societal impact of this work.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
27

---

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We foresee no such risks posed by this work.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: All of the datasets we use are publicly available at huggingface site, and we have properly cited all the training and evaluation datasets used.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paper swi thcode. com/ coomdatasets. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
28

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: This work does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
This work does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This work does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
29


Reviews:
[{"review_id": "upv1ZkocKR", "reviewer": "Reviewer_AbD2", "summary": "The paper explores a new variant of attention with bounded memory to reduce the growing memory size and thus mitigates the memory-intensive challenges of Transformers. The key idea is to set a memory bound, with a predetermined number of usable memory slots, and a gating mechanism to select or mix KV vectors from the previous step.", "questions": "* Fig. 1 can be better illustrated and explained how the proposed Gate Slot Attention works.\n* What is the reason for choosing memory slots to fit in SRAMs? Particularly when the claim is for efficient training. \n* What is the contribution when the memory footprint as shown in Table 2 is larger than other methods? \n* What are the overheads, say additional parameters and operations, as introduced by this work?\n* Some key results could be mentioned in the abstract and introduction.\n* Fix the extra line-breaker as in line 255?\n* The quotation marks used in line 128 to 135 are bizarre.", "limitations": "The limitation mentioned in the paper is moderate.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The paper proposes a variant of attention mechanism with bounded memory, which seems to be an original contribution.", "weaknesses": "* The paper could clarify the significance of its proposed method. The results show not so much improvements compared with prior work in terms of performance and memory costs. The effectiveness and efficiency of the proposed method remain unclear.\n* The writing on the presentation of the proposed method and the evaluation and insightful discussion could be improved.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your thorough review and insightful feedback. \nWe will revise the paper accordingly in the next version to enhance clarity in the abstract, introduction, and discussion sections, polish the writing, fix potential errors, and improve the figures.\n\n---\n\n**Q:** *The paper could clarify the significance of its proposed method. The results show not so much improvements compared with prior work in terms of performance and memory costs. The effectiveness and efficiency of the proposed method remain unclear.\nWhat is the contribution when the memory footprint as shown in Table 2 is larger than other methods?*\n\n**A:** We highlight that Gated Slot Attention (GSA) offers significant advantages over existing methods, as shown in the above tables:\n* GSA shows **improved performance on recall-focused tasks** given limited memory capacity compared with Mamba, RetNet and GLA. Especially on SQUAD and SWDE, GSA outperforms other models by 5.4 and 10.8, respectively. \n* GSA's formulation, which mimics the spikiness in standard attention, facilitates effective continual pretraining from well-trained LLMs like Mistral 7B, demonstrating **superior performance on challenging tasks like NQ, TriviaQA, BBH, and MMLU**.\n\nDespite a slight increase in FLOPs, GSA achieves comparable training throughput to GLA (42K vs. 42.5K tokens/s under 8K context length) and outperforms Flash Attention (36.9K tokens/s) greatly.\nRegarding memory footprint, the increase is minimal compared to GLA for sufficiently large head dimensions (d). The benefits in performance and versatility outweigh this small trade-off.\n\nWe will provide more detailed discussions on complexity and additional results in the next version of our paper to better highlight these contributions.\n\n---\n\n**Q:** *What is the reason for choosing memory slots to fit in SRAMs? Particularly when the claim is for efficient training.*\n\n**A:** GSA involves two passes of gated recurrent processes, necessitating hardware-efficient implementations to match GLA's efficiency.\nTo address this, we follow the approach in [1], accelerating GSA by fully utilizing tensor cores, which is 16× faster than non-matmul counterparts.\n\n[2] reveals that latencies for tensor-core matmuls of 16×16 remain similar to 64×64 (Figure 1), despite operating on 16×16 tiles.\nOur preliminary experiments confirm these findings, motivating us to choose a number of slots that can be processed efficiently in one matmul pass.\n\n| Slots | PPL (↓) | tokens/s |\n| ----- | ------- | -------- |\n| 32    | 13.74   | 46.7K    |\n| 64    | 13.51   | 44.1K    |\n| 128   | 13.46   | 37.1K    |\n\nAs shown in the table, we observe a significant throughput degradation when increasing slots from 64 to 128, while there's a substantial perplexity gap between 32 and 64 slots. \nConsequently, we adopt 64 slots as the default setting, balancing efficiency and performance.\n\nThis hardware-efficient implementation allows GSA to match GLA's efficiency while providing improved performance. \nDetailed implementations are provided in Appendix A. \nWe will include more discussions and comparisons in the next version of our paper.\nThank you.\n\n---\n\n**Q:** *What are the overheads, say additional parameters and operations, as introduced by this work?*\n\n**A:** We appreciate the reviewer's question regarding the overheads introduced by our work.\nWe have included some discussions in Section 3.1 of our paper, but we're happy to provide further clarification here.\n* **Parameter allocation**: For 1.3B parameter models, GSA introduces approximately $dhm \\approx 0.125 d^2$ additional parameters for forget gate mappings compared to Llama.\nThis results in a parameter allocation similar to that of GLA.\n* **Computation Complexity**: The computation complexity of GSA and GLA operation is $O(NCd+Nd^2)$ [3] and $O(NCm+NCd+2Ndm)$, respectively, where N is sequence length, C is chunk size, d is head dimension, and m is number of memory slots. \nWhile we set $m=64$ and $m\\ll d$, the increased complexities are negligible.\n\nOur hardware-efficient implementations demonstrate that GSA performs comparably to GLA in terms of efficiency, as illustrated in Figure 3 in the paper.\nImportantly, this modest increase in parameters and comparable computational efficiency come with significant benefits, which we've discussed in detail in our response to Q1.\nThese advantages justify the minimal overhead introduced by our approach.\n\n[1] FLA: A Triton-Based Library for Hardware-Efficient Implementations of Linear Attention Mechanism: https://github.com/sustcsonglin/flash-linear-attention\n\n[2] Simple linear attention language models balance the recall-throughput tradeoff: https://arxiv.org/abs/2402.18668\n\n[3] Gated Linear Attention Transformers with Hardware-Efficient Training: https://arxiv.org/abs/2312.06635", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nhXAd5qZg7", "reviewer": "Reviewer_LU5r", "summary": "A major challenge is storing information in a bounded number of memory slots. This work builds on ideas from ABC and gated linear attention. ABC recursively updates the bounded-memory keys and values states over time, and computes a softmax with the queries at timestep t to produce the output at t. Gated linear attention and Mamba use data-dependent decays over time to better decide what information to keep versus throw away given the limited memory. GLA, unlike Mamba, enables the use of GPU tensor cores through its parameterization.\n\nGSA retains the selective gating from GLA and the update rules from GSA. The architecture remains chunk-wise parallelizable (like GLA) and provides inductive biases that address limitations of GLA and ABC respectively. GSA is validated up to 2.7Bn parameters from scratch and 7Bn parameters in continual training, providing promising results on standard LM-eval harness benchmark tasks.", "questions": "1. Why do the authors believe the current set of benchmarks is sufficient? Is it possible to include results on retrieval and longer context benchmarks? \n\n2. Can the authors perform error analysis on the continually trained models to help understand why SUPRA and GSA perform poorly on MMLU? \n\n3. Can the authors provide more concrete hypotheses as to where context-aware query tokens (resulting from the first GLA pass) could help in real language modeling settings? \n\nIn the introduction, “slots” are not defined. They are first defined in section 2.", "limitations": "The authors do mention that they ignore retrieval style tasks in their analysis (in Section 6), however if the paper makes claims about memory-efficiency then it is important to evaluate quality on this style of tasks. This is because it is known that there are fundamental memory and quality tradeoffs.\n\nIf the reviewers address these concerns, I will consider increasing my score!", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The writing and contextualization of the contributions with respect to prior work are very clear.\n\nThe architectural inductive baises from GSA versus prior linear recurrent architectures are compelling and very interesting. For e.g., “In GLA the query vector is only dependent on the current input token, while in GSA, the query vector is the output of the first GLA pass and thereby is aware of the entire historical information and in this sense more context aware”. Two passes over the input could help the model make better decisions about what information to keep versus throw away given the limited recurrent memory. \n\nGSA is trained from scratch to relatively large scales (2.7Bn parameters, 100B tokens) and demonstrates high quality on LM-eval harness tasks compared to the baselines, despite using less memory than GLA. \n\nGSA can be implemented efficiently by adopting off-the-shelf algorithms from flash linear attention.", "weaknesses": "Building on my comment on the compelling properties of GSA, the paper does not extensively show how these properties make a difference in language modeling. It would be interesting to know how the “context aware query vectors” help GSA on real-world data by comparing to models without this property. OR for the authors to include discussion on what kind of sequences this property might help with. The architectural modifications are not tied to empirical insights beyond lm-eval harness scores in the current submission.\n\nGSA is only evaluated on LM-eval harness. It is known that models perform somewhat similarly on these short-context tasks with very little memory (Arora et al, 2024) and mentioned in the paper’s limitations section. The paper does not dive into any tasks that require longer context reasoning or retrieval, making it unclear how GSA behaves. \n- \"GSA only needs half the recurrent state size of GLA and a quarter the recurrent size of RetNet, while having better performance, thus enjoying a lower memory footprint during inference\". The claim is not fully validated unless the GSA models are evaluated on tasks that stress-test memory utilization (like retrieval, long-context tasks).\n\nThe MMLU scores of the continually fine-tuned 7B checkpoints remains very low, just like the baseline – the other scores are roughly comparable across models. There is no analysis as to why this is. The paper claims that GSA is drastically better than SUPRA, but the delta is small (<1 point on average), so it is not clear what is meant by this claim.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful feedback and insightful questions.\nWe promise to revise the paper to address your concerns and make the presentation more clear and comprehensive in the next version.\n\n\n---\n\n**Q:** *It would be interesting to know how the \"context aware query vectors\" help GSA on real-world data by comparing to models without this property. \n Why do the authors believe the current set of benchmarks is sufficient? Is it possible to include results on retrieval and longer context benchmarks?*\n\n**A:** Thank you for your very constructive suggestions.\nIn the above table, we have additionally included results to validate the effectiveness of \"context aware query vectors\":\n* We conduct experiments on three tasks proposed by Arora et al, 2024 [1]: FDA, SQUAD, and SWDE, showing **significant improvements in GSA's performance compared to Mamba, RetNet, and GLA on these recall-intensive tasks**. Despite smaller memory capacity, GSA outperforms Mamba, RetNet, and GLA by an average of 5.9, 7.7, and 4.8 points, respectively.\n* We report results on long context benchmarks reported by Xiong et al. 2023 [2]. It is clear that **GSA can be well extrapolated to 16k give 2k training context length**. GSA exhibits better extrapolation than GLA and RetNet on Qasper, NarrativeQA, QuALITY and QMSum. \n\nThese results validate GSA's strong performance on information extraction tasks, which we attribute to more effective memory utilization enabled by \"context aware query vectors.\"\nWe are planning to inspect the performance of GSA scaled to larger model size and number of slots.\n\n---\n**Q:** *The MMLU scores of the continually fine-tuned 7B checkpoints remains very low, just like the baseline – the other scores are roughly comparable across models.\nThe paper claims that GSA is drastically better than SUPRA, but the delta is small (<1 point on average), so it is not clear what is meant by this claim.\nCan the authors perform error analysis on the continually trained models to help understand why SUPRA and GSA perform poorly on MMLU?*\n\n**A:** After carefully diving into the details of SUPRA, we updated the continual pretraining results of GSA models trained with low learning rate of $3×10^{-5}$ combined with slow decay scheduler.\nFrom the above table, we observe that GSA significantly outperforms RetNet and GLA across various tasks, including NQ, TriviaQA, and BBH.\nNotably, for challenging reasoning / language understanding tasks of BBH and MMLU, GSA outperforms GLA by 2.7 and 4.0 points, indicating that GSA's formulation, which is similar to standard attention (SA), allows it to retain more capabilities when finetuned from another SA LLM.\n\nIn the bottom lines we further provide the results of continual pretraining 100B tokens: GSA greatly surpasses SUPRA by 2.2, 5.4, 8.7 and 6.0 points, respectively.\n\n[1] Simple linear attention language models balance the recall-throughput tradeoff: https://arxiv.org/abs/2402.18668\n\n[2] Effective Long-Context Scaling of Foundation Models: https://arxiv.org/abs/2309.16039", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "s4PaPZFRDb", "reviewer": "Reviewer_eYuC", "summary": "The paper introduces Gated Slot Attention (GSA), an enhancement of Gated Linear Attention (GLA), aimed at improving the efficiency of sequence modeling. GSA incorporates a selective gating mechanism to manage memory updates, leveraging a two-pass GLA structure. This approach allows GSA to be more context-aware and to retain a bounded memory footprint, making it suitable for long sequence tasks. \nThe key improvement is the incorporation of a softmax operation to retain sharp attention distributions, which enhances sequence modeling by reducing dilution.", "questions": "How does the inference performance of GSA compare to other SOTA models in terms of speed and memory usage?", "limitations": "Yes", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The GSA's gated update mechanism and two-pass structure offer a straightforward yet effective enhancement over GLA. This design ensures that memory usage remains bounded and manageable, which is critical for handling long sequences efficiently.\n- The experimental results provided by the authors demonstrate the advantages of the GSA mechanism. The results highlight GSA's ability to improve performance in sequence modeling tasks, validating the practical effectiveness of the proposed approach.", "weaknesses": "- While GSA presents improvements, it largely builds on existing GLA techniques. The enhancements, though valuable, might be seen as incremental rather than revolutionary, potentially limiting the perceived impact of the work.\n- Despite efforts to reduce computational overhead, the softmax operation on $QK^T$ still retains a quadratic complexity in training. This raises questions about the authors' claim of linear-time sequence modeling in the title, which could be misleading.\n- The authors do not specify explicitly whether GSA inherits the recurrent architecture of GLA. If this is the case, it would be necessary to compare GSA to other recurrent models like RWKV to provide a comprehensive evaluation of its performance and advantages in sequence modeling tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Q:** *While GSA presents improvements, it largely builds on existing GLA techniques. The\nenhancements, though valuable, might be seen as incremental rather than revolutionary, potentially limiting the perceived impact of the work.*\n\n**A:** We appreciate your feedback and respectfully argue that GSA represents substantial advancements over existing techniques:\n\n* GSA is motivated by viewing $\\mathbf{K}$ and $\\mathbf{V}$ as memories and linearizing standard attention by managing $\\mathbf{K}$ and $\\mathbf{V}$ into constant number of memory slots. This perspective drives our 2-pass recurrence formulation: $\\mathrm{softmax}(\\mathbf{Q}\\tilde{\\mathbf{K}}^\\top)\\tilde{\\mathbf{V}}$. \nThis significantly differ from exising linear attention that conduct linearization by rearranging the computation order of $\\mathbf{Q}\\mathbf{K}^\\top\\mathbf{V}$.\n* The 2-pass recurrence with context-aware queries demonstrates significant improvements in retrieval tasks (FDA, SQUAD, SWDE), showcasing the practical impact of our approach.\n* GSA maintains the general form of standard attention while offering $\\mathbf{K}\\mathbf{V}$ linearization. The $\\mathrm{softmax}$ spikiness [1] notably enhances continual-pretraining results, paving the way for efficient scaling to larger models.\n\n---\n\n**Q:** *Despite efforts to reduce computational overhead, the softmax operation on $\\mathbf{Q}\\mathbf{K}^\\top$ still retains a quadratic complexity in training. This raises questions about the authors' claim of linear-time sequence modeling in the title, which could be misleading.*\n\n**A:** Thank you for your concern! We would like to clarify that in GSA, the $\\mathrm{softmax}$ is applied over $\\mathbf{Q}\\tilde{\\mathbf{K}}^\\top$ rather than $\\mathbf{Q}\\mathbf{K}^\\top$ in standard attention (SA), as presented in Section 2.2.\nUnlike $\\mathbf{K}\\in \\mathbb{R}^{T\\times d}$, $\\tilde{\\mathbf{K}}\\in \\mathbb{R}^{m\\times d}$ is derived by passing $\\mathbf{K}$ along with $\\mathbf{\\phi}\\in \\mathbb{R}^{T\\times m}$ through GLA recurrence, compressing it to a constant size of $m\\times d$.\nThis crucial difference enables GSA to reduce the training complexity from quadratic to linear time, supporting our claim of linear-time sequence modeling.\n\nWe will ensure this distinction is more explicitly stated in the revised version. Thank you.\n\n---\n\n**Q:** *The authors do not specify explicitly whether GSA inherits the recurrent architecture of GLA. If this is the case, it would be necessary to compare GSA to other recurrent models like RWKV to provide a comprehensive evaluation of its performance and advantages in sequence modeling tasks.*\n\n**A:** We would like to clarify that: as detailed in Section 3, GSA's output is indeed derived from $\\tilde{\\mathbf{K}}_t^\\top \\mathbf{q}_t \\in \\mathbb{R}^d$, where $\\tilde{\\mathbf{K}}$ and $\\tilde{\\mathbf{V}}$ maintain constant sizes regardless of sequence length. Both $\\tilde{\\mathbf{K}}$ and $\\tilde{\\mathbf{V}}$ involves one GLA recurrence.\nWe provide hardware-efficient implementations for the 2-pass operations to enhance training efficiency. Full algorithm details are available in Appendix A.\n\nRegarding comparisons with other recurrent models, we include comparisons with RWKV6 in Table 5 in the paper.\nOur results show that GSA significantly outperforms RWKV6 under controlled settings, while RWKV6 is trained on billions of tokens and GSA uses continual pretraining.\nIt's worth noting that while RWKV6 shares similarities with GLA in terms of data-dependent gating and matrix-formed hidden states, it employs additional short convolutions at the expense of efficiency.\n\nDue to limited computational resources, we haven't provided results for RWKV6 models trained from scratch, nor have we employed the specific techniques used by RWKV6, to ensure fair comparisons.\nWe commit to improving the clarity of our presentation in the next version of our paper.\n\n---\n\n**Q:** *How does the inference performance of GSA compare to other SOTA models in terms of speed and memory usage?*\n\n**A:** During inference, GSA demonstrates similar theoretical inference performance compared to other SOTA models :\n* GSA needs $O(1)$ space for hidden states, similar to GLA and other linear models, while SA requires $O(N)$ KV cache.\n* GSA requires $O(1)$ time per step to access memory and generate outputs as in GLA.\n\nOur empirical speed evaluations for 7B models with 2k prefix support these advantages:\n* GSA and GLA consume similar 15GiB memory, while SA requires an additional 3GiB due to KV cache requirements for further generation.\n* Regarding speed, the latency for GSA is 13s, comparable to GLA (13.4s) and SA (12s), aligning with our findings in Tables 2 & 3 of our paper.\n\nIt's important to note that we have not yet fully optimize the recurrent kernel for speeding up the generation stage, nor reducing the I/O overhead through customization.\nWe are committed to exploring these optimizations in future work to fully realize GSA's theoretical inference advantages and enhance its real-world performance.\nThank you.\n\n[1] The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry: https://arxiv.org/abs/2402.04347", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We sincerely thank all reviewers for their thorough examination of our work and their insightful feedback. 
Your thoughtful comments and questions have significantly enhanced our submission and have been addressed in detail in individual responses. 
We have additionally run many new empirical results and analysis, including:
* **Time and memory analysis**: We have added comprehensive efficiency metrics and updated the speed comparisons for different numbers of slots to address Reviewer AbD2's questions regarding computational efficiency.
* **Long-context evaluation**: In response to Reviewer LU5r's concerns, we have performed evaluations on real-world language tasks with extended context, including Qasper, NarrativeQA, QuALITY, and QMSum.
* **Recall-intensive and challenging language reasoning tasks**: To address questions from Reviewers AbD2 and LU5r, we have included results from recall-intensive tasks and more challenging tasks focused on language reasoning to demonstrate the advantages of our GSA approach.

In this shared response, we provide a detailed elaboration of these new results:

### **Recall-intensive tasks**

We evaluated our model on recall-intensive tasks proposed by Arora et al, 2024 [1]:

|            | FDA      | SQUAD    | SWDE     |
| ---------- | -------- | -------- | -------- |
| Mamba      | 9.7      | 33.3     | 34.6     |
| RetNet     | 8.9      | 33.1     | 30.4     |
| GLA        | **11.8** | 25.8     | 43.3     |
| GSA (ours) | 11.3     | **38.7** | **45.4** |

GSA's strong performance, particularly on SQUAD and SWDE, demonstrates that our 2-pass recurrence design enables efficient memory utilization despite limited capacity.

[1] Simple linear attention language models balance the recall-throughput tradeoff: https://arxiv.org/abs/2402.18668

### **Results on challenging language understanding and reasoning tasks**

We expanded our evaluation to include more demanding tasks, following the settings of Llama2 [2] and Mistral:

| Model      | Tokens | NQ       | TriviaQA | BBH      | MMLU     | Avg.     |
| ---------- | ------ | -------- | -------- | -------- | -------- | -------- |
| Mamba      | 1.2T   | 25.4     | 66.2     | 21.5     | 33.2     | 36.5     |
| RWKV6      | 1.4T   | 20.9     | 59.5     | 23.4     | 43.9     | 36.9     |
| SUPRA      | +20B   | -        | -        | -        | 28.0     | -        |
| RetNet     | +20B   | 16.2     | 43.0     | 8.7      | 26.1     | 23.5     |
| GLA        | +20B   | 22.2     | 57.8     | 20.8     | 28.4     | 32.3     |
| GSA (ours) | +20B   | 23.4     | 60.7     | 23.5     | 32.4     | 35.0     |
| SUPRA      | +100B  | 24.7     | 60.4     | 19.8     | 34.1     | 34.7     |
| GSA (ours) | +100B  | **26.9** | **65.8** | **29.3** | **38.1** | **40.0** |

GSA consistently achieves better scores across all tasks compared to SUPRA, RetNet, and GLA at similar token counts and performs competitive with trillion-token models, highlighting its superior in-context learning capabilities and validating our claim that GSA preserves more abilities compared to other linear attention methods under continual-pretraining settings.

[1] Llama 2: Open Foundation and Fine-Tuned Chat Models: https://arxiv.org/abs/2307.09288

### **Length Extrapolation**

In order to showcase the abilities of extrapolating to longer sequences, we conduct additional experiments on long-context tasks, following the settings used in the Llama2-Long report [1].
We tested GSA on four real-world language tasks: Qasper, NarrativeQA, QuALITY, and QMSum. 
The input for each task was truncated at 16k tokens following [1], which is 8x longer than our training context length of 2k tokens. 
The results are presented in the following table:

|        | Qasper   | NarrativeQA | QuALITY  | QMSum    |
| ------ | -------- | ----------- | -------- | -------- |
| Mamba  | 5.6      | **22.2**    | 27.5     | 0.7      |
| RetNet | 11.1     | 0.0         | 26.2     | 0.0      |
| GLA    | 18.4     | 17.2        | 30.9     | 8.9      |
| GSA    | **18.8** | 19.2        | **32.0** | **10.0** |

These results demonstrate that our GSA model generalizes well to sequences of 16k tokens, despite being trained on much shorter contexts. 
Notably, GSA outperforms other linear models across all four tasks, showcasing its robustness and effectiveness in handling long-range dependencies.

[1] Effective Long-Context Scaling of Foundation Models: https://arxiv.org/abs/2309.16039

Meta-review: This paper introduces Gated Slot Attention (GSA) for linear-time sequence modeling, which effectively addresses inefficiencies in traditional attention mechanisms for large language models. All reviewers agree that the method is both novel and impactful. During the rebuttal period, the authors provided additional empirical evidence and clarifications, particularly around memory management and computational performance, which directly addressed the reviewers' concerns. The discussions further solidified that the GSA approach offers substantial advancements over existing methods, validating its effectiveness in a variety of tasks. Based on the strong consensus among reviewers, the thoroughness of the rebuttal, and the high quality of the work, I recommend that the paper be accepted for publication.
Decision: Accept (poster)