Selective Attention Improves Transformer
Authors: Yaniv Leviathan
Matan Kalman
Yossi Matias
Venue: ICLR 2025 Poster (2025)
Keywords: selective attention
attention
transformer
llm
language model
Primary area: foundation or frontier models, including LLMs

Abstract: Unneeded elements in the attention’s context degrade performance. We introduce Selective Attention, a simple parameter-free change to the standard attention mechanism which reduces attention to unneeded elements. Selective attention consistently improves language modeling and downstream task performance in a variety of model sizes and context lengths. For example, transformers trained with the language modeling objective on C4 with selective attention perform language modeling equivalently to standard transformers with ~2X more heads and parameters in their attention modules. Selective attention also allows decreasing the size of the attention’s context buffer, leading to meaningful reductions in the memory and compute requirements during inference. For example, transformers trained on C4 with context sizes of 512, 1,024, and 2,048 need 16X, 25X, and 47X less memory for their attention module, respectively, when equipped with selective attention, as those without selective attention, with the same validation perplexity.

Published as a conference paper at ICLR 2025
SELECTIVE ATTENTION IMPROVES TRANSFORMER
Yaniv Leviathan
Matan Kalman Google Research matanke woannkggoogle.coo
Yossi Matias Google Research yoss i googge com
Google Research
Leviathan@google com
ABSTRACT
Unneeded elements in the attention's context degrade performance. We intro- duce Selective Attention, simple parameter-fiee change to the standard attention mechanism which reduces attention to unneeded elements. Selective attention consistently improves language modeling and downstream task performance in a variety of model sizes and context lengths. For example, transformers trained with the language modeling objective C4 with selective attention perform language modeling equivalently to standard transformers with ~22 more heads and param- eters in their attention modules. Selective attention also allows decreasing the size of the attention's 's context buffer, leading to meaningful reductions in the memory and compute requirements during inference. For example, transformers trained on C4 with context sizes of 512, 1,024, and 2,048 need 16X, 25X, and 47X less memory for their attention module, respectively, when equipped with selective attention, as those without selective attention, with the same validation perplexity.
INTRODUCTION
Different tasks have different memory requirements. On one extreme, copying an arbitrary sequence requires retaining all sequence elements in memory. On the other extreme, determining whether specific element appeared at least once, only requires persisting a constant amount of memory.
Transformers Vaswani et al.) (017)) keep the entire history in their context buffers, allowing them to solve tasks such as copying, while famously leading to their squared attention cost. RNNs (Rumel- hart et al.) (986) and their modern structured state space variants (Gu et al.) [022]] Gu & Dao (024) keep only a constant-sin -ized sketch of the history, making inference cost linear, but rendering them incapable of solving tasks such arbitrary string copying.
Can we design a model that persists just the right amount of context?
Several works (see Section 8) aim to improve costs by compressing otherwise reducing the context size with minimal impact to quality. We take a different approach, focusing instead on quality improvement, and treating cost reductions
side benefit. Specifically, it has been demonstrated (Leviathan] (002) that for some tasks removing unneeded elements from the context buffer enables more efficient transformer programs. Indeed, the attention's 's differentiable memory, all memory cells contribute to the data read, and circuitry is needed to filter out the noise generated by irrelevant memories. Reducing the amount of circuitry needed should improve performance.
In this work we propose Selective Attention, a simple extension to the standard attention mechanism which allows token to decide that another token is no longer needed, reducing the attention that future tokens will pay to it. Selective attention adds no new parameters and only negligible amount of computation, yet yields meaningful improvements synthetic tasks and natural language mod- eling for a range of model and context sizes. Additionally, we show that elements that are selected to be forgotten by selective attention can be safely removed from the attention's is context, leading to substantial reductions in the memory and computation requirements during inference, without pe- nalizing quality. We name our method after the related neuroscience concept of selective attention. Quoting Plebannk & Sloutsky (2017): "Selective attention allows adults to focus task-relevant information, while ignoring task-irrelevant information. This in turn leads to superior processing of task-relevant information."

---

Published as a conference paper at ICLR 2025
Variable Assignment Layer out
1/7
408
238
Parity*
Layer out of 3
Copy Layer
Language Modeling Layer 7 out of 12
piaadort
predeet
Jftrr
Aymee prosdereet
oboma- 136
utorr
<hry
JfM
4ymee-
chan55
Sharcs
XfM
showes
Sharcs
Figure 1: A visualization of the masking by selective attention (red strike-through and attention strength (averaged across heads, blue highlight) for different tasks (see Section 2).
2

---

Published as a conference paper at ICLR 2025
2 MOTIVATING EXAMPLES
Consider a transformer processing an input sequence with three tokens: a, b, C. In a given layer with the standard attention module, token b can decide how much read from token a, and token can decide how much to read from token a, but token b cannot affect how much token e reads from token a. Specifically, if token has determined that token is irrelevant or even misleading to future tokens such as c, there is nothing it can do in the given layer to correct for this. Even in subsequent layers, masking token a is not trivial. Selective attention enables exactly such masking. To illustrate its usefulness, let's consider the Variable Assignment problem, well natural language modeling.
In Variable Assignment the input consists of set of repeated assignments to named vari- ables, followed by a query for the latest value for one of the variables which the model needs to output. For example, for the input: y=7; x=1; x=3; 2=5; x=? the output is 3. Note that Variable Ass Assigt can be seen as a generalization of the Search problem (Leviathan 2022], where repeated occurrences of the query pattern are allowed, and are tasked with finding the most recent occurrence. It is well known that the Search problem is easily solvable by standard transformers, via induction heads (Olsson alll (022). Selective attention facilitates simple reduc- tion from Variable Assignment to Search, whereby every assignment to variable masks out all previous assignments the same variable. In Figure (top) we see that this is indeed the case for transformer trained with selective attention. In Appendix A.1 we show that a transformer with selective attention easily learns general solution to Variable Assignment while a standard transformer does not.
To motivate selective attention for natural language modeling, we first note that Variable Ass i gnmentt is a common sub-task, e.g. when persisting a state. For further motivation, con- sider the common case where a part of the input is ambiguous, and the ambiguity is only resolved at a later token. For example, in the sequence: Bar, ##ack, Obama, the first token Bar encodes several competing meanings, but the later tokens # #ack and Obama resolve it to the entity for the ex-president. For many tasks that mostly concerned with the semantic meaning, later tokens might not want to read the ambiguous meaning from the earlier tokens, so masking them, as with selective attention, might be useful. In Figure [ (bottom) we see that this is indeed the case for transformer trained with selective attention. In the visualized layer, the last token in multi-token expressions masks out the earlier tokens. For example, # #ack masks out bar; obama masks out both bar and # aack; # #bm masks out i; and # i Hl masks out both u and ##c. We also observe additional masking, e.g. the token after masks out the tokens a and day, perhaps because the to- ken a fftee absorbed the meaning from the tokens a and day, or perhaps because the model deems the extra detail is not helpful at this point.
Figure also shows that for the trivial task of Parity', where intermediate results are stored every other token, so that the model's output only a function of the last two tokens, everything but the last two tokens is masked. For the Copy task, selective attention persists the entirety of the string to be copied until copying starts, and then masks out every element as it is copied. See Appendix A.2
3 SELECTIVE ATTENTION
Selective attention is a simple modification on top of standard attention. The key idea is that tokens can mask previous tokens, i.e. the amount of attention a token с pays a previous token a can be reduced by the tokens located between and c. For context size N, we produce a real-valued N N masking matrix S where St, denotes how much token H; masks token 2y (see Section We then constrain S, e.g. to be causal and non-negative (see Section 3.2). We finally accumulate the information in matrix into a new matrix F, taking into account masking by all previous tokens (see Section 3.3 The matrix F = Accumulate( Accumulate(oonstrann (Constrain(S)) is then simply subtracted from the attention logits before applying the softmax:
QKT
SelectiveAAtennionn(..K = softmax
F)V
(1)
Vdl.
Figure illustrates sketch implementation.
3

---

Published as a conference paper at ICLR 2025
3.1 SELECTION FUNCTION
Computing the selection matrix S is akin to computing a compatibility score, and many functions be used, for example, separate bilinear form. Following observation from Leviathan (2022) that a common case for a token wanting to mask another is after absorbing its contents (i.e. after attending to it). we instead simply reuse the result of one of the existing heads of the attention module. This means that the selection function adds no new parameters or computation. Note that the head still contributes to attention as usual. We also experimented with separate bilinear form, and in spite of adding additional parameters and computation, this resulted in the same slightly worse results (see Appendix A.3)
3.2 CONSTRAINTS
Following observations from Leviathan (0022 apply the following constraints S:
1. Zero out negative values (i.e. applying ReLU), only reducing attention, never boosting it.
2. Zero out the first column, so as not to mask the <BOS> token.
3. Zero out the diagonal, so as not to let token mask itself.
We observe that all three constraints improve performance (see Appendix A.5).
3.3 ACCUMULATION
In this work focus transformer decoders, selective attention cannot influence the atten- tion operation by past tokens. We chose cumulative summation as our accumulation function. We observe some improvement by only applying the masking for future tokens (i.e. the masking by to- ken would not affect its own attention operation), so Fi,j = Eksi-1 Sk.j (see ablation Appendix A.4
attn logits einsum ("bhnd, bhmd->bhnm", Q, K) sqrt (dk)
attn logits = where (causal mask, attn logits, float ("-inf"))
s = attn__too [etn__oggtts:; 0]
# Select head 0.
(S)
# Only positive selection.
s(..., 0]
# # Do not mask self.
mask <BOS>.
(1 eye *
s = roll(S. (S, 1, -2); S(..., 0, :3::]
# Mask strictly in the future.
mp. cumsum (S, axis=-2)
# Accumulate.
attn logits Fl;, None]
attn weights softmax (attn logits)
Figure 2: A sketch implementation of selective attention. The colored lines the additions to standard attention.
CONTEXT PRUNING
As presented in Section 3 while beneficial model quality, selective attention has negligible impact on inference efficiency However, an additional modification can improve inference efficiency substantially. Specifically, selective attention can reduce the memory and computation requirements of the attention module, via pruning elements from the context buffer.
To how, note that once token is sufficiently masked by selective attention, it will not contribute meaningfully to any future attention operations. Such tokens can be safely evicted from the context buffer. We could pick fixed threshold, and prune all elements whose soft masking higher than the threshold (i.e. Fij > T), but that would make the memory savings hard to take advantage of (e.g.
TThe additional computation of (hn?²) is negligible compared the ((d)²) of standard attention.

---

Published as a conference paper at ICLR 2025
due to fragmentation and variable number of dropped tokens each iteration). Instead, we observe that the sparsity (i.e. the magnitude of the masking) per layer is stable across samples (see Section Z). To that effect we set fixed memory budget for each layer, which directly translates memory and compute savings. Since when a token is dropped it remains dropped for all future tokens, given memory budget of K = K1, K1,-.., Kr tokens in layer I we drop nothing, and for each following token we drop the single not-yet-dropped past token with the highest F value. This maintains no more than K₁ tokens layer 1. Given an overall memory budget, we allocate it between the layers, using greedy iterative approach. For context size N, we initialize K N...., . KΣ = N. In each iteration t, we set K K for all n F m, and K Km - C for constant C' (we use 8 in our experiments), where m = arginin, f(- C.), where C. = (O......0, 0, C, .,....0) contains at the vth position. In other words, we iteratively reduce the memory budget of the layer that impacts model performance the least. We stop when model performance reaches a predefined threshold, in our experiments, the performance of standard transformer without selective attention.
tokens for each layer, when processing the first K
Note that with a low memory budget there might be some discrepancy between training and infer- ence. Fine tuning the model after the budgets have been set (or even better, in each iteration) might be advantageous and lead to larger reductions in memory budgets, but we haven' 't experimented with this setup yet.
As selective attention's 's masking is beneficial for the model, we observe significant reductions in context sizes without any auxiliary losses (see Section 6.2p. Nevertheless, we can further encourage the model to mask out more elements by adding an explicit term to the loss, like so:
EL_ max; M Lmem = L + I, -a
Where L is the standard loss llog-perplexity in the case of language modeling), c is a small weight factor (in our experiments we set = 0.1 without further tuning), L is the number of layers, pad is the number of non-pad tokens, and M = i - Ei= min (in(F",, T//T is our approximation for the memory requirements at the ith token for layer 1 (0- < Mf < +). We clamp Ft from above by not to reward increasing it indefinitely (I'ss already clamped from below to 0). We set t 7= without further tuning. Since the memory required for a given layer is the maximum memory required for each of the tokens, the loss only considers the maximum among the Mís. We observe further reduction in context sizes with this explicit loss term (see Section 6.2
EXPERIMENTAL SETUP
In all of our experiments we use decoder-only transformer with multi-head attention, as in Vaswani jet al. (2177), with the following modifications: we use Pre-LN instead of Post-LN (Xiong et al. 2020 learned position encoding, SwiGLU gates instead of MLPs (Shazeer 2020). normalize the Q and K projections (Dehghani et al., 2023 remove the biases (Raffel et al. 2023). and replace the LayenNorms with RMSNorm (Zhang & Sennrich 2019). Note that we tested other variants, including a vanilla decoder-only transformer exactly as in (Radford et al. 2019) and observed similar results. We trained our models with the Adam optimizer with B1 = 0.9 and 88 = 0.999 for total of 524,288 steps. We used cosine decay and 1,000 linear warmup steps and learning rate of 0.005. We repeated some of the experiments with different learning rates and obtained similar results. We used a batch size of 256 and a fixed context size of 512 for all training runs except for the context size experiments (Figure left) where we used batch of 128. We follow Esser et al. (2024), and parameterize model size by parameter d such that dmode! 64d and Theads = Wlayyers d (see Table X in Appendix 13.1) We trained all of our models on TPUv4s. For the language modeling experiments, we used the C4 (Raffel et al.] (203) dataset with vocabulary of size 8K tokens built with the SentencePiece tokenizer (Kudo & Richardson 2018). We repeated some of the experiments with vocabulary of size 32K and observed similar results. We also ran experiments with WikiText (Merty et al.) 2016), and Imlb (Chelba et al. 2014) and observed similar results.
For larger models, we observed more cases of divergence when not normalizing the Q and K projections.
5

---

Published as a conference paper at ICLR 2025
RESULTS
6.1
GENERATION QUALITY
Transformers with selective attention perform consistently better, measured by perplexity the validation set, across model and context sizes. We also observe consistent improvements on a set of downstream tasks.
Figure 3 (left) compares the validation perplexity for causal language modeling on the C4 dataset of decoder-only d 12 transformer models with and without selective attention, for varying context sizes, Likewise Figure 3 (right) compares validation perplexity with and without selective attention, for varying model sizes with context length of 512. We observe improvements across model sizes, and that the improvements grow with the size of the context.
2.825
128
d=8
With Selectiva Attention
2.9
Without Selective Attontion
2,800
Valication log-erppexxin
256
2,7
d-12
2.750
2.6
d=14
2.725
258
d-16
1024
2048
d=18
d=20
2.675
2.4
d=22
d=24
dw2n
2.650
2048
d=
500
1000
1500
2000
0.0
0.2
0.4
0.6
0.8
1.0
Context length
parameters
1e9
Figure 3: (Left) The validation perplexity of ad = 12 transformer, with (blue) and without (orange) selective attention, for varying context sizes. (Right) The validation perplexity of transformers of various sizes, with (blue) and without (orange) selective attention, for context size of 512.
Figure 4 shows that even when equipped with additional attention heads (and increasing the param- eters of the attention module proportionally, so that the size of each head remains constant), trans- formers with standard attention only become comparable to those with selective attention, when they have about double the number heads (and parameters) as their selective attention counterparts.
2m88
2010
hel2
he34
d-10
2.3
X
18
di:12
(o)) octioon altnl
d4:14
Volcation bo perpleeity
ded setective attention
2.0
2012 hol4 1o10
d-10 wi select attention
26-11
U-12 selective attentiu
attn))
select
held pel6pe10
selective
X6
(W) selective adtr
2.8
Jmid 1n
20 20
selective
Ae 16 10u:2 20u 24
2.5
X-16
(W)) selective attal
Attontion
Figure 4: Perplexity of transformers of various sizes with and without selective attention. For the cases without selective attention we add additional attention heads with their respective parameters (i.e. increase the sizes of all projection matrices). Transformers with selective attention perform equivalently to those with standard attention with ~2X as many heads and parameters.

---

Published as a conference paper at ICLR 2025
model accuracy the ARC (Clark et al. 2018), HellaSwag (Zellers et al. 2019), PiQA (Bisk et al. 2019). CommonSenseQA (Talmor et al, 2019), and OpenBookQA (Mihaylov et al. 2018) benchmarks, with and without selective We observe consistent improvements across all
In addition to perplexity the validation set, also
attention, for models various sizes (see Table model sizes with selective attention.
Table 1:
Accuracy numbers for baseline transformers with standard attention and transformers with selective attention for various model sizes and downstream tasks.
ARC (Easy)
ARC (Challenge)
HellaSwag
Base
Selective
Base
Selective
Base
Selective
16
38.46%
39.33%
23.24%
23.24%
38.83%
40.01 %
d = 18
40.48%
41.10%
24.67%
24.86%
41.77%
42.60 %
20
41.33%
41.85%
25.14%
25.52%
43.97%
45.08 %
22
42.43%
42.70%
25.83%
26.41%
46.04%
47.71%
= 24
43.83%
44.47%
26.64%
26.76%
48.70%
50.32% %
26
44.70%
45.58%
26.37%
26.95%
51.02%
52.25% %
28
45.64%
46.95%
26.76%
27.37%
53.50%
53.76%
SAnMen
OpenBookQA
PiQA
Base
Selective
Base
Selective
Base
Selective
16
25.44%
25.82%
34.26%
34.33%
68.07%
68.49 %
18
26.45%
26.69%
35.32%
35.42%
69.05%
69.88 %
= 20
26.30%
26.63%
35.42%
35.27%
69.82%
70.49%
22
26.56%
27.10%
35.96%
36.88%
70.55%
71.43%
24
27.49%
27.49%
36.36%
37.13%
71.34%
71.95%
26
27.70%
27.78%
37.64%
37.52%
71.89%
72.47%
= 28
27.86%
28.51% %
37.33%
37.54%
72.32%
72.85 %
6.2 INFERENCE EFFICIENCY
Efficiency improvements via selective attention stem from a reduction the attention module's context size when using pruning as in Section 4 Specifically, a smaller context translates directly to more efficient inference common scenarios. Indeed, note that during inference with a large context and batch size (bn >> d), loading the KV-cache (linear in the size of the context) dominates the memory bandwidth (Pope et al. 2022), which is often the bottleneck for generation (Shazeer, 2019). In addition, for very large context sizes (n >> d). taking the dot product of the query and the keys in the cache and calculating the weighted average of the values in the cache both dominate compute, i.e., in this setup, smaller context directly translates similar gains in FL.OPs.
When pruning the context with selective attention, we measure substantial improvements in the memory requirements for the attention module, at the same or better perplexity than the baseline without selective attention. Figure in Appendix A.6 illustrates the trade-off between perplexity and efficiency.
For example, for d = 12 transformers with context sizes of 512, 1,024, and 2,048, we see that with selective attention we can maintain the baseline's perplexity while reducing the memory re- quirements of the attention module by factors of 5X, 7X, and 8X respectively, without any explicit losses (i.e. using only the standard language modeling objective). When training with the L.nem loss (Equation and an value of 0. the improvements grow to 16X, 25X, and 47X respectively. We also measure the memory savings when considering only very long examples by filtering C4 to only include examples that are at least 90% the size of the context buffer. In that settings get memory savings of 12X, 18X, and 24X, while maintaining the perplexity of the baseline without selective attention. To maintain the perplexity gains of selective attention, instead of matching the perplexity the baseline (i.e. the rightmost point on the flat part of the blue graphs of Figure 6). we need 3X, 4X and 4X less memory, for the context sizes above, respectively.

---

Published as a conference paper at ICLR 2025
In all cases above, we optimized the per-layer budgets training set and reported results separate unseen test set.
Finally, we compare selective attention to other attention variants and pruning methods and observe that it achieves a substantially better quality-cost tradc-off than all baselines (see Appendix A.100).
SELECTION PATTERNS
It is interesting to question which elements from the context are being masked by selective attention for language modeling. Figure illustrates the values of the F matrix for a specific example (see Appendix A.7 for the full example text). We observe that some layers (e.g. 6) are mostly dense (i.e. low values), while other layers (e.g. 2) are sparse (i.e. high F values). As expected, all layers persist some of the most recent elements, but several of the sparse layers (e.g. layers 1, 4, and 9) also persist elements for long time periods, as can be seen by the vertical lines. This suggests that simply limiting the attention module a local window would result in the same quality gains as those achieved by selective attention, which we confirm in Appendix A.8
Layer
Layer
Layer
Leyer
Layer
Layer
Loyer 6
Loyer? 7
Layor
Layer 9
Layor 10
Layer 11
Figure 5: Visualization of the F matrix (greener is lower, i.c. less masking) for a d 12 transformer for the text in Appendix A.7
Figure 9 Appendix A. illustrates the values of the F matrix averaged over 1,000 examples, and demonstrates that the sparsity patterns are stable across examples. Interestingly, we observe that these sparsity patterns are sometimes stable across different training runs, hinting at some general properties of language modeling (see Figure T0 in Appendix (A.1)...
Figure I Appendix A. depicts the items remaining in the context buffer after pruning (see Section 4). with budget set to match the perplexity of standard transformer. We observe that the per-layer budgets correspond well to the values of the F' matrix as seen in Figures 5 and 9 For example, layer which gets the highest budget also has the lowest F values. This might indicate that the scales of the F values are consistent across layers.
Figure 12 in Appendix illusttatss the value of the last row of F (i.e. the masking for the 512th token) for each of the layers. We observe some interesting patterns, for example, layer persists the end-of-sentence periods (*.).
8

---

Published as a conference paper at ICLR 2025
8 RELATED WORKS
Transformer Improvements. Our work aims improving the transformer architecture. Since its introduction in Vaswani et al. (2017), a large volume of research proposed architecture modifications towards an improved model. Some notable works here include Pre-LN instead of Post-LN (Xiong
et al. 2020), gated units instead of MLPs (Shazeer, 22020), removing the biases (Raffel et al. 2023),
using RMSNorm instead of LayerNorm (Zhang & Sennrich 2019), normalizing the Q and K pro-
jections (Dehghani et al. 2023), and multi-query and group-query attention (Shazeer) 2019 Ainslie
ct al. 2023).
Attention Modifications. Our work focuses on modifying the attention module. Most of the ex- isting research work around modifying attention focuses on a different goal than ours, specifically on devising more efficient attention variants. Some such variants are based on approximations to the attention operations and include sparse attention approximations (Child et al., 2019; Ding et al.) 2023), and linear attention approximations (Shen et al., 2024; Katharopoulos et al., 2020] Schlag et al. 2021). Some works focus on hardware-aware optimizations instead, such as FlashAttention
(Dao et al. 2022) and Ring Attention (Liu et al. 2023a).
Context Pruning. A part of our work (Section consists of removing elements from the context buffer. Several works employ this mechanism order to increase inference efficiency. Among those are compression methods, that aim to learn a compressed representation for tokens in the context buffer, with without auxiliary losses, in order to replace several tokens with their compressed form (Munkhdalai et al., 2024; Ren et al. 2023 Yun et al.] 20233; Mu et al., [024] Rae et al, 2019). Another line of work tries to simply remove elements from the context buffer without replacing them with new compressed forms, with minimal negative impact to model quality. The simplest and most widely adopted of these is using local attention windows in some of the layers (Wang et al.) 2019 BBlttgy et al.) 2020 Zaheer et al., (02II). More sophisticated variants employ heuristics t0 evict less useful tokens from the context buffer instead of just the earliest ones (Oren et al. 2024 Zhang et al. 2023 Liu et al. 2023b Berchansky et al. 20233 Ge et al. 2024). Anagnostidii et al. 2024) fine tunes existing models to learn to prune tokens from the context buffer with similarities to selective attention, but is more involved (e.g. needs root solving for evaluating the a-sigmoid), produces binary prune decisions, and notably doesn't improve quality.
Inference Efficiency. A part of our work (Section(4 4 consists of making inference from transformers more efficient. Many approaches aim to speed up inference from transformers, including distilla- tion (Hinton al., 2015)). sparsification (Jaszczur et al. (02T), quantization (Hubara et al. 2016).
architecture modification (So et al., 2022; Shazeer 2019). and algorithmic optimization (Dao et al.
2022
Leviathan et al. 2022).
Finally, we note that the importance of learning to forget has been shown repeatedly in many works, more generally beyond transformers. One of many notable examples are the forget-gates in LSTMs (Hochreiter & Schmidhuber 1997).
DISCUSSION
In this work we introduced Selective Attention, simple parameter-free change to the standard at- tention mechanism which consistently improves language modeling performance across model sizes and context lengths, and can lead to substantial inference efficiency improvements. Given that it adds no new parameters, only negligible amount of compute, and provides consistent improvements, selective attention might be good default for transformer decoders.
Future directions. We applied selective attention to decoder-only transformers. It could be inter- esting to investigate its applicability to encoders as well (see Appendix A.9 for initial results in an encoder-decoder setup); Reducing the size of the context as in Section 4 ciency but not training efficiency. It might be interesting to explore iteratively reducing the size of the context buffer during training; We did not further train the models after removing elements as per Section It seems conceivable that further improvements could be achieved with some additional training after context reduction; We only experimented with pre-training models with selective atten- tion. is interesting to investigate how it could be applied in a fine-tuning step to existing models; While observed similar results with selective attention in several setups (Section 5). there
improves inference effi-

---

Published as a conference paper at ICLR 2025
still important variants did not test, notably transformers with multi-query (Shazeer 2019) and grouped-quers query (Ainshi et al., (023) attention, as well as models much larger than 1B parameters; Finally, selective attention can be implemented in GPU-aware way, similar to Flash Attention (Dao et al. 2022).
10 IMPROVING NEURAL ARCHITECTURES
In The Art of Transformer Programming, Leviathan 2022) selected a set of foundational problems (sorting, searching, addition, etc.) and manually implemented transformers to solve them (i.e. by manually setting the model's weights). They showed that several programs become much easier, especially for small transformers, when equipped with mechanism allowing selectively mask items in the context buffer, similar to selective attention. They further hypothesized that such mechanism will have similar positive effects on language modeling, which motivated our work.
Zhou et al. 2023) proposed the generalize on task the task can solved by a short RASP program which works for all input lengths", i.e. that problems that are easily solved by transformers are those that are easily solved by human programmers using RASP. It follows that problems that are not easily solved by humans us- ing RASP are hard for transformers well, and if we made those easier, by changing the transformer architecture (and respectively the capabilities of RASP) we could meaningfully improve transform- Similarly, when constructing transformer programs by hand, Leviathan (022) notes that "... the most interesting cases are those that are hard for us humans and are hard for the optimizer or the architecture, and understanding these better might be key to creating better Al systems."
Conjecture, that "Transformers tend to length
We are strong advocates of this method, and believe that finding basic problems for which we can- not program general solution by hand neural model, fertile approach for architecture improvements.
ACKNOWLEDGMENTS
We'd like to extend a huge thank you to Raya Leviathan, Blake Hechtman, Asaf Aharoni, Uri Mendlovic, Danny Lumen, Avinatan Hassidim, Dani Valevski, Eyal Segalis, Molad Eyal, the Theta Labs and Google Research teams, and our families for insightful feedback, ideas, suggestions, and support.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Joshua Ainslie, James I Lee-Thorp. Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head check- points, 2023. URL URLhttps:// http:s://aaxiv..oo /wwgg/abs/// 205. 13245.
Sotiris Anagnostidis, Dario Pavllo, Luca Biggio, Lorenzo Noci, Aurelien Lucchi, and Thomas Hof- mann. Dynamic context pruming for efficient and interpretable autoregressive transformers, 2024.
URL https: org/
15805
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long- transformer, 2020. URLpttts::/ http://aaxvvvv.om orq/ abs/2004... 05150
Moshe Berchansky, Peter Izsak, Avi Caciularu, Ido Dagan, and Moshe Wasserblat. Optimizing retrieval-auggentte reader models via token elimination, 2023. /ttssttss///axxi
abs/2310.13622
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language, 2019.
11641
Ciprian Chelba, Tomas Mikolov, Mike Schuster; Qi Ge, Thorsten Brants, Phillipp Koehn, and Tony Robinson. One billion word benchmark for measuring progress in statistical language modeling, 2014. /ttpsss:://ax http:////aa.ii..../a/a//////////
3005
Rewon Child, Scott Gray, Alee Radford, and llya Sutskever Generating long sequences with sparse transformers, 2019. htttp:://axxxiiiii
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge, 2018. Bttpttts: /tttps://arxi..... http/vxxviv.//// org/ abs/ awtpaaas...ae/ee 00057754777777
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness, 2022. URL https:/t::/// wwg. abs/ 22C5.
Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, Rodolphe Jenatton, Lucas Beyer, Michael Tschannen, Anurag Amab, Xiao Wang. Carlos Riquelme, Matthias Min- derer, Joan Puigcerver, Utku Evci, Manoj Kumar, Sjoerd van Steenkiste, Gamaleldin F. Elsayed, Aravindh Mahendran, Fisher Yu, Avital Oliver, Fantine Huot, Jasmijn Bastings, Mark Patrick Collier, Alexey Gritsenko, Vighnesh Birodkar, Cristina Vasconcelos, Yi Tay. Thomas Mensink, Alexander Kolesnikov, Filip Pavetić, Dustin Tran, Thomas Kipf, Mario Lučié, Xiaohua Zhai, Daniel Keysers, Jeremiah Harmsen, and Neil Houlsby. Scaling vision transformers to 22 billion parameters, 2023. URL wttphhtpst... prtts:://axiv. org/ abs/ /a.bb/23333 22422
Jiayu Ding, Shuming Ma, Li Dong, Xingxing Zhang, Shaohan Huang, Wenhui Wang, Nanning Zheng, and Furu Wei. Longnet: Scaling transformers to 1,000,000,000 tokens, 2023. URL
https://arxivs attp:::/arxxx..orr/////////333 wwtp://ggg/b//eooe
02486
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, Dustin Podell, Tim Dockhorn, Zion En- glish, Kyle Lacey, Alex Goodwin, Yannik Marek, and Robin Rombach. Scaling rectified flow transformers for high-resolution image synthesis, 2024. URL http://axiv..ro//aaa
Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive kv cache compression for 1lms, 2024. URLhttts://axxiv org/
abs/23100 01801
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, 2024.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces, 2022.
11

---

Published as a conference paper at ICLR 2025
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilliing the knowledge in neural network, 2015.
URL httppss://axxi...rrtt abs/1503. 02531
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8):
1735- 1735-1780, 1997.
Itay Hubara, Matthieu Courbariaux, Daniel Soudry, Ran EI- El-Yaniv, and Yoshua Bengio. Quantized neural networks: Training neural networks with low precision weights and activations, 2016. URL
attps://arxiis wtggg/a//////0 07066
Sebastian Jaszczur, Aakanksha Chowdhery, Afroz Mohiuddin, Lukasz Kaiser, Wojciech Gajewski, Henryk Michalewski, and Jonni Kanerva. Sparse enough in scaling transformers, 2021. URL
https://arxiv.
2111. 12763
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers
mns: Fast autoregressive transformers with linear attention, 2020. URL https://arxix:
arg/abs/200 2006 16236
Taku Kudo and John Richardson. Sentencepiece: simple and language independent subword tokenizer and detokenizer for neural text processing, 2018.
1808 06226
Yaniv Leviathan. The art of transformer programming, 2022. github. io/taotp. html
URL https://yanivle.
Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding, 2022. URL htttttttps/aaaaett http:///aaaxxvvooo///////1//1// wttp
171992
Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near- infinite context, 2023a. URL bltttss:/aaxxiv.rgggg
01889
Zichang Liu, Aditya Desai, Fangshuo Liao. Weitao Wang. Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, and Anshumali Shrivastava. Scissorhands: Exploiting the persistence of importance hypothesis for I1m kv cache compression at test time, 2023b. URL aattps://arxiv..
abs faas/2305.v 17118 17118
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering, 2018. URL httpps://axxxii.
org/abs/ 1809, 02789
Jesse Mu, Xiang Lisa Li, and Noah Goodman. Learning compress prompts with gist tokens, 2024.
URL http:httts//axx... org abs/ /255//30002 2304 08467]
Tsendsuren Munkhdalai, Manaal Faruqui, and Siddharth Gopal. Leave no context behind: Efficient infinite context transformers with infini-attention, 2024. URL
http://axvvvvooo....e /
24C4.07143
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-oodds Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads, 2022. URL http::ttt///aaaa httpppss////x..../////a org/ abs / 220//229... 11895
Matanel Oren, Michael Hassid, Nir Yarden, Yossi Adi, and Roy Schwartz. Transformers are multi- state rnns, 2024. URL
06104
Daniel J. Plebanek and Vladimir M. Sloutsky. Costs of selective attention: When children no- tice what adults miss. Psychological Science, 28:723 732, 2017. URL https: : //pp.. semant eschol ar. httpa//aar.ooooooooaaaa133333010
12

---

Published as a conference paper at ICLR 2025
Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Anselm Lev- skaya, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. Efficiently scaling trans- former inference, 2022. Uttptts:: http::://axxvvvvvv...
05102
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Lan- guage models are unsupervised multitask learners, 2019. URL https::/ http::::/cdd..peeeee openai.. Comn///eteer--agggggee-oooeeee Languaag models are_aareee____ uns.pervised. Seitittask_Lezzersss
Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar, and Timothy P. Lillicrap. Compressive transformers for long-range sequence modelling, 2019. URL https://arxiv.ont http://aaxxx.oooaa////
1911.05507
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer, 2023. URL http:p/s//aaa....og/aa//////s//s
10683
Siyu Ren, Qi Jia, and Kenny Q. Zhu. Context compression for auto-regressive transformers with sentinel tokens, 2023. URL httppss//aaxiivoorraaaas//////// 0ttp//1112100.1.....2222
Adam Roberts, Hyung Won Chung, Anselm Levskaya, Gaurav Mishra, James Bradbury, Daniel Andor, Sharan Narang, Brian Lester, Colin Gaffney, Afroz Mohiuddin, Curtis Hawthorne, Aitor Lewkowyez, Alex Salcianu, Mare van Zee, Jacob Austin, Sebastian Goodman, Livio Baldini Soares, Haitang Hu, Sasha Tsvyashchenko, Aakanksha Chowdhery, Jasmijn Bastings, Jannis Bu- lian, Xavier Garcia, Jianmo Ni, Andrew Chen, Kathleen AoneatlyJJnaahen H. Clark, Stephan Lee, Dan Garrette, James Lee- Thorp. Colin Raffel, Noam Shazeer, Marvin Ritter, Maarten Bosma, Alexandre Passos, Jeremy Kaitin-Shepard Noah Fiedel, Mark Omernick, Brennan Saeta, Ryan Sepassi, Alexander Spiridonov, Joshua Newlan, and Andrea Gesmundo. Scaling up models and data with t5x and seqio, 2022. URL http:ttss:
David E. Rumelhart, Geoffrey E. Hinton, and J. Williams. Learning internal representatiin by error propagation, 1986. URL https: /app. semanticscholasnt http:h/oo.......ooooooss/s/sssss
62245742
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers, 2021. URL
11174]
Noam Shazeer. Fast transformer decoding: One write-head all you need, 2019. URL https:
7/arxiv. //a.xxv.org/abs//////oo/o
Noam Shazeer. Glu variants improve transformer, 2020. URL https:://axiv.. www.iiv.orggae www.orgg........o....... /bs/
2002. 05202]
Zhuoran Shen, Mingyuan Zhang, Haiyu Zhao, Shuai Yi, and Hongsheng Li. Efficient attention: Attention with linear complexities, 2024. URL
01243
David R. So, Wojeiech Mańke, Hanxiao Liu, Zihang Dai, Noam Shazeer, and Quoe V. Le. Primer:
Searching for efficient transformers for language modeling, 2022. URL https://arxiv:
erq/abs/2100 08668
Alon Talmor, Jonathan Herzig, Nicholas Lourie, and Jonathan Berant. Commonsensequ: A question answering challenge targeting commonsense knowledge, 2019. URL https://azxiv.ort
abs/1811.000 00937
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, 2017.
Zhiguo Wang, Patrick Ng, Xiaofei Ma, Ramesh Nallapati, and Bing Xiang. Multi-passage BERT: A globally normalized BERT model for open-domain question answering. In Kentaro Inui, Jing Jiang, Vincent Ng, and Xiaojun Wan (eds.), Proceedings of the 2019 Conference Empir- ical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pp. 5878-5882, Hong Kong, China, Novem- ber 2019. Association for Computational Linguistics. doi: 10.18653/v//119-5550 URL https:: http:///c..anttoooooooo///////-.
13

---

Published as a conference paper at ICLR 2025
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tie- Yan Liu, On layer normalization in the transformer architec- ture, 2020. /tttth////axiiii org/ abs 220222202...2 04745
Jungmin Yun, Mihyeon Kim, and Youngbin Kim. Focus on the core: Efficient attention via pruned token compression for document classification. In Findings of the Association for Computational Linguistics: EMNLP 2023. Association for Computational Linguistics, 2023. doi: 10. 1e.188//1/2223..nmmm// Lindings-emnlp.900 URL http://áss doi. 10o..ooo//.....55////10 2023. indinng-emnn 909
Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula. Qifan Wang, Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences, 2021. URLhttps: //tp://.aa.xi./////a//aa/aamaaav org/ /2bs/2011 14062
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a ma- chine really finish your sentence?, 2019. URL httphhttt://axiiis
Biao Zhang and Rico Sennrich. Root mean square layer normalization, 2019. URL https:// arxiv.orgaabs///ooooooo
07467)
Zhenyu Zhang. Ying Sheng. Tianyi Zhou, Tianlong Chen, Lianmin Zheng. Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, Zhangyang Wang, and Beidi Chen. H2o: Heavy- hitter oracle for efficient generative inference of large language models, 2023.
Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Josh Susskind, Samy Bengio, and Preetum Nakkiran. What algorithms can transformers learn? study in length generalization, 2023. URL httpt:tt////mmm /tatts://arxvy ozggabs/2310.1100/000
14

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 TRANSFORMERS WITH SELECTIVE ATTENTION LEARN A GENERAL SOLUTION TO VARIABLE ASSIGNMENT
Transformers with selective attention reach close to O validation loss and 100% precision extremely fast when trained on Variable Assignmen ansigttent and they generalize well to out of distribution cases, unlike transformers without selective attention.
Setup. We train small transformers (d 3), with and without selective attention, to solve the Variable Assi gnment problem with 3 variables, 1,000 possible values, and 128 assignments. We train with a batch size of 2,048 for 65,536 steps.
In distribution. The transformer with selective attention reaches validation loss of 0.002 (and 100% accuracy) after less than 1,000 training steps. The transformer without selective attention only achieves validation log- peepperplextty loss of 3.18 and 26% accuracy after 1,000 step. At the end of training (65,536 steps) the transformer with selective attention obtains a loss of 2.2e-8, whereas the transformer with standard attention is at 0.01. Both transformers reach 100% accuracy at the end of training.
Out of distribution. We observe much stronger generalization capabilities for the transformer with selective attention. When we run on an out distribution test set with the same 3 variables but only 2 possible values, the transformer with standard attention's 's accuracy drops substantially to 70% (loss of 3.64). Meanwhile the transformer with selective attention maintains 100% accuracy (with a loss of 2.4-
We observed similar results in other settings (e.g. 10 variables and 10 possible values). We also repeated the experiments with somewhat larger transformers = 8) and observed similar results.
See Appendix A.lI for an example of the attention patterns for the Variable Assignment task.
A.2 ADDITIONAL SYNTHETIC TASKS
We validate selective attention two additional synthetic tasks, Copy and Parity', that are the two opposite extremes in terms memory requirements. See Figure]
Setup. We train small transformers (d 3), with and without selective attention, solve the Copy and Parity" tasks. We train with a batch size of 2,048 for 65,536 steps.
Copy. Here the transformer gets arbitrary sequence of varying length delimited by special token and tasked with outputting a copy of the sequence. We used lengths that are uniformly distributed between and 24. The context size for this task is 3 x Limaxs where Imaz 24 is the length of the longest possible input sequence (the 3 extra tokens are for <BOS>, <EOS> and the special aoddof-inputtsequenne token). To solve this task, the model cannot forget anything before copying starts, after which point it can mask out tokens that were already copied.
Parity". Here the transformer gets a binary sequence where bits in the odd positions random and bits in the even positions contain the parity of all bits in the earlier odd positions. Equiva- lently, bits in the even positions contain the XOR of the bits in the previous two positions, so the two previous positions are enough for computing the next token. The loss only considers the even positions.
Unsurprisingly, transformers with selective attention achieve practically 0 loss (and 100% accuracy) the end of training, as do standard transformers without selective attention, on both tasks.
15

---

Published as a conference paper at ICLR 2025
A.3 SEPARATE BILINEAR FORM
We experiment with using a separate bilinear form instead of reusing the output of an existing attention head. We compare transformers (for d= 8 and d 12) trained with selective attention C4 for 524,288 steps, to similarly trained transformers where the selection function is implemented a separate bilinear form (adding additional parameters and computation). The transformers with standard selective attention (i.c. sharing the outputs of an existing attention head) achieve the same or slightly better log-pews pee-perllexitis the end of training (see Table
Table The log-perplexity on the validation set after 524,288 training steps for (1) standard atten- tion, (2) selective attention with a separate bilinear form for the selection module (more parameters than the baseline), and (3) selective attention.
8
Additional parameters
Standard attention
2.96
2.68
No
Selective attention (separate)
2.91
2.63
Yes
Selective attention
2.90
2.63
No
A.4 SELF-IMPACT
Table compares the results of forbidding self-impact (i.e. not allowing token affect its own attention operation, as in Section 3.3) to those when allowing it (i.e. not shifting the matrix S). As can be seen, the shifting provides small but consistent improvement.
Table 3:
The average log-perplexity on the validation set of training runs after 65,536 training steps for selective attention vs selective attention without shifting, for various model sizes.
d = 10
d
12
d 14
d
16
d
18
d 26
Selective attention (no shift)
2.927
2.832
2.753
2.692
2.641 2.639
2.516
Selective attention
2.923
2.831
2.750
2.691
2.511
16

---

Published as a conference paper at ICLR 2025
A.5 ABLATING THE CONSTRAINTS
We ablate the constraints selective attention applies to S (see Section 3.2
Negative selection. While with selective attention token decide to reduce attention another token by all future attention operations, allowing a token to strengthen another token's ' contribu- tion all future attention operations does not make sense. Indeed, when removing this constraint (dropping the ReLU, so that S can contain negative values) the training does not converge.
Masking the <BOS> token. Since several algorithms can benefit from the existence of the sentinel <BOS> token (Leviathan, 2022). it is plausible that masking it is detrimental. When we allow selective attention mask the <BOS> token, we observe neutral slightly worse results compared to standard selective attention where the <BOS> token is forced to never be selected, see Table
Self-masking. Since selective attention reuses an existing attention head the selection function, motivated by the absorption observation (see Section 3.1 it seems plausible that token should never mask itself. Indeed, when we allow tokens mask themselves (i.e. we stop zeroing the diagonal of 5') we observe worse results, see Table
Table 4: The log- perplexity on the validation set after 524,288 training steps for (1) selective atten- tion without the <BOS> constraint and (2) selective attention.
d
Selective Attention (w/o <BOS> constraint)
Selective Attention
12
2.6409
2.6373
14
2.5486
2.5483
16
2.4750
2.4741
18
2.4153
2.4156
20
2.3674
2.3673
24
2.2909
2.2865
Table 5: The log-perplexity the validation set after 524,288 training steps for (1) selective atten- tion without the self-masking constraint and (2) selective attention.
d Selective Attention (w/o self masking constraint)
Selective Attention
12
2.7348
2.7251
14
2.6510
2.6423
18
2.5261
2.5209
17

---

Published as a conference paper at ICLR 2025
A.6 PERPLEXITY-EFFIIENN TRADE-OFF
Figure illustrates the trade-off between perplexity gains and efficiency gains when pruning as in Section See Section for details.
Context size 512
2.72
Valication log perpleeexxty
2.70
2.68
2.60
Expi
Crem
10
20
30
40
50
Context size 1,024
2.72
Validation log-perplextty
2.69
2.66
Capl
2.64
Crem
10
20
30
40
50
Context size 2,048
2.72
validation log-nerplexity
2.68
2.66
Capi
2.64
Crem
10
20
30
40
SU
Memory improvement factor
Figure 6: The trade-off between perplexity and KV-cache size for d = 12 transformers with context sizes of 512, 1,024, and 2,048. Note that in all cases the perplexity with selective attention is better equal to that of the baseline without selective attention (the dotted lines). Selective attention trans- formers trained with the Lmem loss and e = 0.1 (Equation 2 match the perplexity of the baseline with 16X, 25X, and 47X less memory, while those with the standard loss match the perplexity of the baseline with 5X, 7X, and 8X less memory, respectively.
18

---

Published as a conference paper at ICLR 2025
A.7 EXAMPLE DETAILS
The following text from the C4 validation set was used in Figures 5 and 12
the real problem with eraditional dental veneers has little to do who determines nes what normal aesthetically pleasing smile looks like is the real issue america struggled for decades with defining image a marketplace bent on exploiting peoples flaws for economic gain the danger of this national obsession has become systemic since the days of twiggy fashion magazines offer photoshopped perfection the standard to which we should aspire the effects of this insidious marketing made their way into breast implants and the definition of hollywood smile carving their bodies and their teeth people use their resources to chase a false picture of their perfect self people make investments in the tens of thousands of dollars at their dentist office to get the perfect smile this message has become so endemic that pperfect hollywood
with how they function or their performance
people with nice smiles are convinced only one particularly relevant example of this
smile is acceptable is highlighted in a june 2015 article entitled saving jane gary nankin dds discusses how he " saved the smile of patient who was not content with her first set of # porcelain veneers
smile
endodontic referral for treatment of tooth number 15 , followed by composite core build up 2 periodontal therapy in both the anterior region and upper left to achieve optimal tissue health 4 preparation of maxillary teeth and placement of permanent restorations of dental implant by the periodontist followed by preparation of mandibular teeth and placement of permanent restorations fully healed and ntegraatee implant in the
5 placement
restore the pesition of tooth number 30 regarding person s smile , the strong link self esteem and self worth make an imperfect set of teeth concern however the picture in the article clearly illustrates what appears to be a well constructed and healthy - looking smile the entire premise is puzzling how does a dentist promote " saving " a smile that 97 $ of the people in america would love to show off ?
19

---

Published as a conference paper at ICLR 2025
A.8 COMPARISON WITH LOCAL ATTENTION
Table 6 compares the validation perplexity of = 12 transformers with various local attention patterns (all-local and alternating), that of standard transformer and to that of transformer with selective attention. For the all-local attention transformers we set all layers to be sliding window attention layers with a fixed sized window. For example. "All-local 32" denotes transformer where all tokens can only attend up to 32 tokens back. We also include transformers with alternating local and global layers, where we have local attention layers followed by global attention layer, in repeated fashion. For example, "Local-globe 32" denotes a transformer with 3 local attention layers where tokens can only attend up to 32 tokens back, followed by global layer where tokens can attend to all past tokens, and this 4-layer structure is repeated for the 12 layers of the d 12 transformer. We report the perplexity numbers after 524,288 training steps. We observe that all local attention patterns perform worse than the dense baseline, which in turn performs worse than transformer with selective attention.
Table 6: Validation log-perplexity of transformers with different local attention patterns.
Model Type
Validation Log-Perplexins
All-local 32
2.7860
All-local 64
2.7386
All-local 128
2.7154
All-local 256
2.6981
All-local 384
2.6873
All-local 448
2.6849
All-local 480
2.6834
Local-global 32
2,7046
Local-global 64
2,7105
Local-global 128
2.7154
Local-global 256
2.6993
Local-global 384
2,6895
Local-global 448
2,6870
Local-global 480
2.6861
Baseline (standard attention)
2.6815
Selective attention
2.6372
A.9 RESULTS ON T5
We experimented with applying selective attention while pre-training T5 (Raffel et al. 2023). Here we use the standard T5 pre-training recipe and code from T5X (Roberts et al.) 2022). Specifically, in this setup, the model is encoder-decoder, and we apply selective attention to the decoder only, leaving the encoder as-is. The standard T5 span-corruption pre- wre-training objective used, as in Raffel al. (2023), both for training and the reported metries. See Table for the results.
We observe some improvements for T5 with selective attention for the tested model sizes.
Table 7: The span corruption loss per non-padding token on the validation set of training runs after 524,288 training steps for baseline T5 encoder-decoder vs T5 encoder-decoder where the decoder equipped with selective attention, for various model sizes.
TS-small TS-base
T5-large
T5
1.962 1.952
1.693
1.528
T5 with selective attention
1.691
1.522
20

---

Published as a conference paper at ICLR 2025
A.10 COMPARISON WITH EFFICIENT ATTENTION METHODS
Figure 7 compares selective attention to other efficient attention and attention pruning methods, including HzO (Zhang et 2023), TOVA (Oren al., 2024). and sparse attention (Child al. 2019). We also include Window + following Oren et al. (2024). Note that H2O, TOVA, and Window can be applied post-training, whereas selective attention and sparse attention require training the mode
We observe that the tested setting of language modeling C4, selective attention substantially outperforms all of the tested efficient attention baselines.
Context size 512
3.00
Vanilla attention
2.95
Selective attn
Validation log-perpiee lexitty
Selective attn Cyaam
TOVA
2.85
Window+4
2.80
H2Q
Sparse atin
2.75
2.70
2.65
2.60
Context size 1,024
3.00
Vanilla attention
2.95
Selective aitn
Idation log oog-errau
2.90
Selective attn Cxem
TOVA
2.85
Window+4
2.80
H20
Sparse attn
2.75
2.70
2.65
1.00
1.25
150
1.75
2.00
2.25
2.50
2.75
3.00
Context size 2,048
Vanilla attention
2.95
Selective attn
Validatio Validtioon log-peralevty
2.90
Selective attn
TOVA
2.85
Window
2.80
Sparse attn
2.75
2.70
2.65
2.60
LU0
1.25
150
1.75
2.00
2.25
2.50
2.75
3.00
Memory improvement factor
Figure 7: The trade-off between validation perplexity and KV-cache size for d 12 transformers with context sizes 512, 1,024, and 2,048 for selective attention and other efficient attention mech- anisms.
Also Appendix A.8 for comparison with models trained with various local attention patterns.
21

---

Published as a conference paper at ICLR 2025
A.11 EXAMPLE ATTENTION PATTERNS FOR VARIABLE ASSIGNMENT
Figure X shows an example of the attention patterns for all 3 layers of transformers with and without selective attention, for the Vari able Assignment task. The part of the example sequence shown starts with the tokens: 2=, 177, Y=, 661, z=, 114, z=, 468.
For the first layer of the transformer with selective attention we observe that tokens of type "<value>" (e.g. 177, 6 61, etc., i.c., those at the even positions) attend to themselves and to the immediately preceding token (of type "<variable>=" c.g. 2-, Y-, etc.). A potential explanation that this allows the value tokens to absorb the information of the variable they are assigned to, so that from this point onwards, the token would contain representation of the pair (variable, value), and the preceding variable-only token would no longer be relevant for future tokens. E.g., the token 177 might contain a combined representation of 2 and 1 177. In this same layer, tokens type "<variable>=" attend to themselves.
We further observe that for the transformer with selective attention, the attention patterns in the second and third layers are almost identical. In these layers, tokens in the even positions, now containing a combined (variable, value) representation according to the postulate above, attend to themselves. Tokens in the odd positions, still containing representation of the assigned variable according to the hypothesis above, attend to the value of the last assignment to the same variable. This allows masking the previous assignment (see Figure T.
In contrast, the baseline transformer without selective attention (which doesn't It solve the Vari ab e Ass i gnment. task for out of distribution problems, see Appendix I. exhibits attention patterns that are harder to interpret.
Transformer (w/ Selective Attention)
Transformer (w/ Vanilla Attention)
Layer
Layar
116
891
2-1
116
24
881
3M 2
116
2
2x
NM
Layer 2
Layet
Layer
100
27
110
Figure 8: The attention patterns (attention strength averaged across heads) for all layers of trans- formers with selective attention (left) and with vanilla attention (right).
22

---

Published as a conference paper at ICLR 2025
A.12 ADDITIONAL FIGURES FOR CONTEXT PRUNING
Figure 9 iillustrates the values of the matrix averaged across 1,000 examples of length at least 512 from the C4 dataset.
Figure comparrss the values from Figure those obtained from a different training run (different random initialization and different data shuffle). While this isn't always the case, sometimes observe stable sparsity patterns like those in this example, hinting at some general properties of language modeling C4.
Figure I illustrates the tokens that remain in the context buffer after pruning (as in Section for the example text (see Appendix A.7 for d 12 model with selective attention, trained with the Lmem loss (Equation 2 0.1), for a memory budget where the validation perplexity matches transformer without selective attention. The per-layer memory budgets chosen by the pruning algo- rithm for this model are: [B, 48, B, 8, 24, 8, 168, 16, 8, 64, 8, 8], leading to a memory saving factor of 16X.
Figure 12 shows which tokens are pruned for the example text in Appendix A.7
O
Layer
Layer
Layer
Layer 4
Layer
Layer
Leyer 7
Layer 8
Layer
Layer
Figure 9: Visualization of the F matrix (greener is lower, i.e. less masking) for a d = 12 transformer averaged across 1,000 examples.
23

---

Published as a conference paper at ICLR 2025
Layerd
Lowe
-We11
Liperd
Loyar
Lype 12
Lewe11
Figure 10: Visualization of the F matrix (greener/bluer is lower, i.e. less masking) for d 12 transformer averaged across 1,000 examples for two training runs (different random initialization, and different shuffle of the training data). While only have anecdotal evidence, it is interesting that we sometimes observe these stable sparsity patterns across training runs.
Layer 0
Layer
Layer
Layer
Layer
Layor
Layer
Layer
Layer 8
Layor 9
Layer 10
Layer 11
Figure 11: Visualization of the persisted elements for d 12 transformer for the text in Appendix
A.7
The white pixels denote tokens removed from the context buffer per Section 4
24

---

Published as a conference paper at ICLR 2025
Laye
vene carta
poreease com
Figure 12: A visualizing of the elements that are masked for the last 1512hh token, for a d = 12 transformer for the text in Appendix A.7 We observe some interesting patterns, for example, layer 4 persists the end-of-sentence periods
25

---

Published as a conference paper at ICLR 2025
A.13 PARAMETER COUNTS
Table 8 shows the actual parameter counts for models with different ds (see Section 5
Note that
selective attention does not add any extra parameters.
Table 8: The number of model parameters for different values of d as in Section S
d Playrs
nheade
dmodel
Number parameters
8
8
512
33,603,584
10
10
10
640
59,699, 200
12
12
12
768
97,615,877
14
14
14
896
149,666, 048
16
16
16
1024
218,226,66 688
18
18
18
1152
305,717,760
20
20
20
1280
414,387,200
22
22
22
1408
546.641,
24
24
24
1536
704,950,272
26
26
26
1664
891, 468,032
28
28
28
1792
L,108,655.88
26


Reviews:
[{"review_id": "AnVLc5er9k", "reviewer": "Reviewer_L9tT", "summary": "The paper introduces a parameter-free modification to the transformer attention mechanism called Selective Attention, which reduces the attention to irrelevant or unneeded tokens in a sequence. This approach enhances performance in language modeling tasks by reducing memory usage and computational costs without compromising model quality. Experiments show that transformers with Selective Attention achieve comparable performance to larger models with traditional attention mechanisms, providing improved efficiency and effectiveness across various model sizes and tasks.", "questions": "1. Have you considered evaluating Selective Attention on other architectures, such as encoder or encoder-decoder transformers? If so, could you share any preliminary results or insights?\n\n2. How sensitive is the performance of Selective Attention to the memory budget per layer? Could you offer guidelines on tuning this budget for tasks beyond language modeling?\n\n3. Have you tested Selective Attention on tasks other than language modeling, such as question answering or text summarization? How might the approach handle tasks with varied context and attention needs?\n\n4. Could you provide further comparisons between Selective Attention and other efficient attention methods, such as sparse or linear attention, to clarify the unique advantages of Selective Attention?\n\n5. How does context pruning with Selective Attention affect tasks requiring long-range dependencies? Are certain types of tokens or information more prone to being pruned?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. **Efficient Memory Management**: The Selective Attention mechanism effectively prunes unneeded tokens, significantly reducing memory usage during inference without degrading model performance. This efficiency gain is particularly valuable for scaling transformers in resource-constrained environments.\n2. **No Additional Parameters**: Selective Attention operates without introducing new parameters or significantly increasing computational overhead, which preserves the simplicity of the transformer architecture.\n3. **Theoretical Motivation and Experimental Rigor**: The paper provides a thorough theoretical motivation for selective attention, supporting it with extensive experimental results on benchmarks like HellaSwag and C4, which demonstrate the effectiveness of the approach.", "weaknesses": "1. **Limited Scope of Model Architectures**: The experiments are primarily conducted on decoder-only transformer models. Further analysis is needed to verify if Selective Attention can similarly benefit encoder models or encoder-decoder models used in other tasks, such as translation or summarization.\n2. **Potential Over-Reliance on Hyperparameter Tuning**: Selective Attention’s performance may depend on optimal memory budget settings per layer, which could complicate deployment in different tasks or models. Although a memory reduction process is described, further tuning could be required in practical applications.\n3. **Dataset and Task Diversity**: While Selective Attention shows improvements in language modeling tasks, testing it on a wider range of tasks (e.g., text generation or long document understanding) would strengthen the case for its generalizability and adaptability to diverse applications.\n4. **Comparison with Similar Methods**: The paper could benefit from a more detailed comparison with other recent efficient attention mechanisms, such as sparse attention or adaptive attention approaches, to highlight any relative advantages Selective Attention may have.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "adavC4hGO0", "reviewer": "Reviewer_5FJ3", "summary": "This paper introduces *Selective Attention*, a parameter-free modification to the standard attention mechanism in transformers which helps reduce the attention given to irrelevant elements in the context. To do so, they introduce a *selection mask* in the attention computation to mask “irrelevant” token, and propose using the previous tokens’ attention on other tokens (on a specific head) as masking for future attention computations. On language modelling on C4, results show that transformers equipped with Selective Attention can achieve comparable performance to standard transformers with twice the attention heads and parameters, and slightly better performance on the downstream task of HellaSwag.\n\nThe authors then propose a method that uses the selection mask to prune elements from the attention's context buffer, reducing memory and computational demands during inference, and show that it can lead to significant memory saving (up to 47x for large context sizes) while preserving performance of non-selective attention transformers.", "questions": "- In the selection attention computation, why use a single head’s logits to mask all heads? Did you try keeping each head separate (i.e. using the logits head i for the selection mask of head i)?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 3, "strengths": "- Given the simplicity of the proposed approach and the reported performance gains, it seems that selective attention could be a significant addition to the transformer architecture if it is further validated.\n- Even without the performance gains, the efficiency gains (particularly without having to modify the pretraining loss) are quite relevant and make Selective Attention look like a viable alternative to other efficient attention mechanisms.", "weaknesses": "- The main problem of this paper is weak experimental validation. The authors only show gains in the language modelling task (using the relatively noisy and deprecated C4 dataset) and on a single downstream task (where they show smaller gains). They don't compare to existing pretrained models or other efficient attention mechanisms. While the (limited) experimental results are promising, they are not enough to validate the approach. I suggest replicating the recipe of an existing (state-of-the-art) pretrained LLM, replacing the attention mechanism with Selective Attention. Additionally, in the current LLM area, more downstream tasks need to be tested to validate the approach. For the efficiency section, the authors should compare with other efficient attention mechanisms.\n- The presentation paper could also be better: it uses a whole page for a single figure and further page for intuition, both of which could be much shorter. It then moves decently important plots to the appendix. I suggest restructure the paper to give more space to the experimental results and comparison to other methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "NrvsuVE6vY", "reviewer": "Reviewer_T9JB", "summary": "This paper argues that, for transformer-based language models, selectively attending to a subset of the tokens in the context can improve the model's quality. Motivated by this intuition, the paper proposes a new attention technique, where if a token has already received high attention weights by previous tokens, it means that its content has already been \"absorbed\" in the model's representations, and therefore it should receive less attention weights onward. The paper proposes an implementation for this idea by reusing the attention logits of one of the attention heads, and thereby adding no additional parameters and only a small amount of computational overhead. Further, the proposed method also allows for pruning some tokens from the context when they will never be attended to again, which leads to memory saving.\n\nExperiments with a 100M transformer based LM trained on the C4 dataset suggests that the proposed approach achieves promising results on language modeling perplexity and one downstream task (HellaSwag).", "questions": "None", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The motivation is clear, and the execution of the idea to not to tokens that are already attended to is clever and clean\n- The proposed method is easy to implement, adds no parameters, and saves memory overhead\n- Interesting results and analysis", "weaknesses": "- My main concern is about the limited empirical evaluation where only one downstream task is considered. As the paper argues, \"different tasks have different requirements,\" it is crucial to explore whether selective attention is broadly applicable by evaluating it on a diverse set of tasks with different requirements. This can be complemented with synthetic evaluations that might require the model to store the entire sequence, e.g., counting the number of a certain token in a sequence. \n- Besides a more comprehensive evaluation, I also encourage the authors to strengthen their findings by trying selective attention on larger-scale models and datasets", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "w3yRMUPvbt", "reviewer": "Reviewer_obGb", "summary": "This paper presents selective attention, a method for adjusting the attention weight of tokens based on the amount of attention they received from previous tokens. The authors also propose to use this method to drop some of the history tokens with the lowest (adjusted) weight, also proposing a new loss term to transformer-based LLM training. Extensive experiments show that this method both consistently improves performance, and reduces memory requirements substantially.", "questions": "The authors mention that they compute the S matrix using \"one of the existing heads\" (#167). Their code indicates that they simple take the first one. Is this ablated in any way? Does taking a different head / an average of all heads make much difference?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- The method is simple and intuitive\n- Results are very strong and consistent, showing both accuracy improvements, and very impressive memory reduction\n- Experiments are quite extensive\n- Given these results, overall I tend to agree with the authors' statement in Section 9: \" Given that it adds no new parameters, only a negligible amount of compute, and provides consistent improvements, *selective attention might be a good default for transformer decoders*.\"", "weaknesses": "While I like this paper a lot, I have several reservations.\n\nFirst, the level of novelty is not huge. The idea of pruning the least important token in each iteration is similar in spirit to cache-pruning methods such as H20 (Zhang et al., 2023), TOVA (Oren et al., 2024) and others. The authors basically apply a similar approach at training time. While this method works great (as noted above), which is a sufficient contribution on its own, and also present some novel concepts (e.g., the new loss term), it would be better to tone down the novelty claims a bit.\n\nSecond, I am not sure I am fully convinced by some of the major claims in the paper.\n \n(a) the examples in Section 2 are very appealing, but they are missing a critical baseline: how much attention does the vanilla transformer assign the pruned tokens in such cases? Perhaps it already knows it needs to ignore them? Again, the method works great so obviously something good is happening, but it is not clear the intuition is capturing the essence of it.\n\n(b) similarly, the authors say (#113) \"if token b has determined that token a is irrelevant or even misleading to future tokens such as c, there is nothing it can do in the given layer to correct for this.\". But I am not sure how their method allows the transfer of such negative signal. If I understand correctly, the authors penalize tokens that received high attention in previous steps (I.e., their value in F is high). I am not sure how information about irrelevant or misleading tokens can be propagated in such cases. \n\nThird, while the writing is generally clear, I found section 3 a bit confusing, requiring several reads to understand a fairly simple method. For instance, it is never explicitly mentioned that each token is penalized by the amount of attention it got from all subsequent tokens. I would also advise the authors to formally define masking, as it is not entirely trivial in this context.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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:

Decision: Accept (Poster)