Selective Attention Improves Transformer

## Abstract

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 improves language modeling performance in a variety of model sizes and context lengths.
For example, a range of transformers trained with the language modeling objective on C4 with selective attention perform equivalently to standard transformers with ∼ similar-to \sim ∼ 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 with 100M parameters 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.

## 1 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 a specific element appeared at least once, only requires persisting a constant amount of memory.

Transformers (Vaswani et al., ) 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 (Rumelhart et al., ) and their modern structured state space variants (Gu et al., ; Gu & Dao, ) keep only a constant-sized sketch of the history, making inference cost linear, but rendering them incapable of solving tasks such as 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 or otherwise reducing the context size with minimal impact to quality.
We take a different approach, focusing instead on quality improvement, and treating cost reductions as a side benefit.
Specifically, it has been demonstrated (Leviathan, ) that for some tasks removing unneeded elements from the context buffer enables more efficient transformer programs.
Indeed, in the attention’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 a 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 a negligible amount of computation,
yet yields meaningful improvements in synthetic tasks and natural language modeling 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 context,
leading to substantial reductions in the memory and computation requirements during inference, without penalizing quality.
We name our method after the related neuroscience concept of selective attention. Quoting Plebanek & Sloutsky (): “Selective attention allows adults to focus on task-relevant information, while ignoring task-irrelevant information. This in turn leads to superior processing of task-relevant information.”

Figure: Figure 1: A visualization of the soft-masking by selective attention (red strike-through) and attention strength (averaged across heads, blue highlight) for different tasks (see Section 2).
Refer to caption: x1.png

## 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 to read from token a, and token c can decide how much to read from token a, but *token b cannot affect how much token c reads from token a*. Specifically, 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.
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, as well as natural language modeling.

In Variable Assignment the input consists of a set of repeated assignments to named variables, 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; z=5; x=? the output is 3.
Note that Variable Assignment can be seen as a generalization of the Search problem (Leviathan, ), where repeated occurrences of the query pattern are allowed, and we 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 et al., ).
Selective attention facilitates a simple reduction from Variable Assignment to Search, whereby every assignment to a variable masks out all previous assignments to the same variable.
In Figure 1 (top) we see that this is indeed the case for a transformer trained with selective attention.
In Appendix A.1 we show that a transformer with selective attention easily learns a general solution to Variable Assignment while a standard transformer does not.

To motivate selective attention for natural language modeling, we first note that Variable Assignment is a common sub-task, e.g. when persisting a state.
For further motivation, consider 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 are 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 1 (bottom) we see that this is indeed the case for a 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 ##ack; ##bm masks out i; and ##la 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 token after absorbed the meaning from the tokens a and day, or perhaps because the model deems the extra detail is not helpful at this point.

Finally, Figure 1 also shows that for the trivial task of Parity^*, where intermediate results are stored every other token,
so that the model’s output is 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.

## 3 Selective Attention

Selective attention is a simple modification on top of standard attention.
For context size $N$, we produce a real-valued $N\times N$ soft-mask matrix $S$ where $S_{i,j}$ denotes how much token $x_{i}$ masks token $x_{j}$ (see Section 3.1).
We then constrain $S$, e.g. to be causal and non-negative (see Section 3.2).
We finally accumulate the information in matrix $S$ into a new matrix $F$, taking into account masking by *all* previous tokens (see Section 3.3).
The matrix $F=\text{Accumulate}(\text{Constrain}(S))$ is then simply subtracted from the attention logits before applying the softmax:

$$ $\text{SelectiveAttention}(Q,K,V)=\text{softmax}(\frac{QK^{T}}{\sqrt{d_{k}}}-F) V\\$ (1) $$

Figure 2 illustrates a sketch implementation.

### 3.1 Selection Function

Computing the selection matrix $S$ is akin to computing a compatibility score, and many functions can be used, for example, a separate bilinear form.
Following an observation from Leviathan () 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 a separate bilinear form, and in spite of adding additional parameters and computation, this resulted in the same or slightly worse results (see Appendix A.2).

### 3.2 Constraints

Following observations from Leviathan (), we apply the following constraints to $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 a token mask itself.

We observe that all three constraints improve performance (see Appendix A.3).

### 3.3 Accumulation

In this work we focus on transformer decoders, so selective attention cannot influence the attention 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 a token would not affect its own attention operation), so $F_{i,j}=\sum_{k\leq i-1}S_{k,j}$ (see ablation in Appendix A.4).

Figure: Figure 2: A sketch implementation of selective attention. The colored lines are the additions to standard attention. Note that selective attention adds a negligible amount of computation, however care should be taken in practice in order to not interfere with compiler optimizations.

## 4 Context Pruning

As presented in Section 3, while beneficial to model quality, selective attention has negligible impact on inference efficiency(^1^11The additional computation of $\mathcal{O}(hn^{2})$ is negligible compared to the $\mathcal{O}(dn^{2})$ of standard attention.).
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 see how, note that once a 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 a fixed threshold, and prune all elements whose soft masking is higher than the threshold (i.e. $F_{i,j}>\tau$), but that would make the memory savings hard to take advantage of (e.g. due to fragmentation and a variable number of dropped tokens each iteration).
Instead, we observe that the sparsity (i.e. the magnitude of the soft-masking) per layer is stable across samples (see Section 7).
To that effect we set a fixed memory budget for each layer, which directly translates to memory and compute savings.
Since when a token is dropped it remains dropped for all future tokens,
given a memory budget of $K=K_{1},\dots,K_{L}$ tokens for each layer,
when processing the first $K_{l}$ tokens in layer $l$ 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_{l}$ tokens in layer $l$.
Given an overall memory budget, we allocate it between the layers, using a greedy iterative approach.
For context size $N$, we initialize $K^{0}_{1}=N,\dots,K^{0}_{L}=N$.
In each iteration $t$, we set $K^{t}_{n}=K^{t-1}_{n}$ for all $n\neq m$, and $K^{t}_{m}=K^{t-1}_{m}-C$ for a constant $C$ (we use 8 in our experiments), where $m=\operatorname{argmin}_{i}PPL(\cdot|K^{t-1}-C_{i})$, where $C_{i}=(0,\dots,0,C,0,\dots,0)$ contains $C$ at the $i$th 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 a standard transformer without selective attention.

Note that with a low memory budget there might be some discrepancy between training and inference. 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 masking is beneficial for the model, we observe significant reductions in context sizes without any auxiliary losses (see Section 6.2).
Nevertheless, we can further encourage the model to mask out more elements by adding an explicit term to the loss, like so:

$$ $\mathcal{L}_{mem}=\mathcal{L}_{ppl}+\epsilon\cdot\frac{\sum_{l=1}^{L}\max_{i}M ^{l}_{i}}{{L\cdot n_{\neq pad}}}$ (2) $$

Where
$\mathcal{L}_{ppl}$ is the standard log-perplexity loss,
$\epsilon$ is a small weight factor (in our experiments we set $\epsilon=0.1$ without further tuning),
$L$ is the number of layers,
$n_{\neq pad}$ is the number of non-pad tokens,
and $M^{l}_{i}=i-\sum_{k=1}^{i}\min(F^{l}_{i,k},\tau)/\tau$ is our approximation for the memory requirements at the $i$th token for layer $l$ ($0\leq M^{l}_{i}\leq i$).
We clamp $F^{l}_{i,k}$ from above by $\tau$ so as not to reward increasing it indefinitely ($F$ is already clamped from below to 0). We set $\tau=1$ 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_{i}^{l}$s.
We observe further reduction in context sizes with this explicit loss term (see Section 6.2).

## 5 Experimental Setup

In all of our experiments we use a decoder-only transformer with multi-head attention, as in Vaswani et al. (), with the following modifications: we use Pre-LN instead of Post-LN (Xiong et al., ), learned position encoding, SwiGLU gates instead of MLPs (Shazeer, ), normalize the Q and K projections (Dehghani et al., )(^2^22For larger models, we observed more cases of divergence when not normalizing the Q and K projections.), remove the biases (Raffel et al., ), and replace the LayerNorms with RMSNorm (Zhang & Sennrich, ).
Note that we tested other variants, including a vanilla decoder-only transformer exactly as in (Radford et al., ) and observed similar results.
We trained our models with the AdamW optimizer with $\beta_{1}=0.9$ and $\beta_{2}=0.999$ for a total of 524,288 steps. We used cosine decay and 1,000 linear warmup steps and a 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 3 left) where we used a batch of 128.
We follow Esser et al. (), and parameterize a model size by a parameter $d$ such that $d_{model}=64d$ and $n_{heads}=n_{layers}=d$ (see Table 7 in Appendix A.9).
We trained all of our models on TPUv4s.
For the language modeling experiments, we used the C4 (Raffel et al., ) dataset with a vocabulary of size 8K tokens built with the SentencePiece tokenizer (Kudo & Richardson, ).
We repeated some of the experiments with a vocabulary of size 32K and observed similar results. We also ran experiments with WikiText (Merity et al., ), and lm1b (Chelba et al., ) and observed similar results.

## 6 Results

### 6.1 Generation Quality

Transformers with selective attention perform consistently better, as measured by perplexity on the validation set, across model and context sizes. We also observe consistent improvement on a downstream task, HellaSwag (Zellers et al., ).

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 a context length of 512. We observe improvements across model sizes, and that the improvements grow with the size of the context.

Figure: Figure 3: *(Left)* The validation perplexity of a $d=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 a context size of 512.
Refer to caption: extracted/5899480/images/comb.png

Figure 4 shows that
even when equipped with additional attention heads (and increasing the parameters of the attention module proportionally, so that the size of each head remains constant), transformers with standard attention only become comparable to those with selective attention, when they have about double the number of heads (and parameters) as their selective attention counterparts.

Figure: Figure 4: Perplexity on the C4 validation set after 524,288 training steps of transformers with various sizes (parameterized by $d$ as per Section 5) with and without selective attention. For the cases without selective attention we add additional attention heads with their respective parameters (i.e. we increase the sizes of all projection matrices). Transformers with selective attention perform equivalently to those with standard attention modules with $\sim$2X as many heads and parameters.
Refer to caption: extracted/5899480/images/extra_heads.png

In addition to perplexity on the validation set, we also measure model accuracy on the HellaSwag benchmark (Zellers et al., ), with and without selective attention, for models of various sizes (Table 1). We observe consistent improvements across all model sizes with selective attention.

**Table 1: Comparing the accuracy of transformers of different sizes with and without selective attention on the HellaSwag benchmark.**
|  | Accuracy (baseline) | Accuracy (selective attention) |
| --- | --- | --- |
| $d=16$ | 38.8% | 40.0% |
| $d=18$ | 41.8% | 42.6% |
| $d=20$ | 44.0% | 45.1% |
| $d=22$ | 46.0% | 47.7% |
| $d=24$ | 48.7% | 50.3% |
| $d=26$ | 51.0% | 52.2% |
| $d=28$ | 53.5% | 53.8% |

### 6.2 Inference Efficiency

Efficiency improvements via selective attention stem from a reduction in the attention module’s context size when using pruning as in Section 4.
Specifically, a smaller context translates directly to more efficient inference in 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., ), which is often the bottleneck for generation (Shazeer, ).
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, a smaller context directly translates to similar gains in FLOPs.

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 6 in Appendix A.5 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 requirements 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 $\mathcal{L}_{mem}$ loss (Equation 2) and an $\epsilon$ value of 0.1, 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 we get memory savings of 12X, 18X, and 24X, while maintaining the perplexity of the baseline without selective attention.
Finally, to maintain the perplexity gains of selective attention, instead of matching the perplexity of the baseline (i.e. taking 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.
In all cases we optimized the per-layer budgets on a training set, while the reported results are on a separate unseen test set.

## 7 Selection Patterns

It is interesting to question which elements from the context are being masked by selective attention for language modeling.
Figure 5 illustrates the values of the $F$ matrix for a specific example (see Appendix A.6 for the full example text).
We observe that some layers (e.g. 6) are mostly dense (i.e. low $F$ 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 to a local window would not result in the same quality gains as those achieved by selective attention, which we confirm in Appendix A.7.

Figure 7 from Appendix A.8 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 8 in Appendix A.8).

Figure 9 from Appendix A.8 depicts the items remaining in the context buffer after pruning (see Section 4), with a budget set to match the perplexity of a standard transformer.
We observe that the per-layer budgets correspond well to the values of the $F$ matrix as seen in Figures 5 and 7. For example, layer 6 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 10 in Appendix A.8 illustrates 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 4 persists the end-of-sentence periods (‘.’).

Figure: Figure 5: Visualization of the $F$ matrix (greener is lower, i.e. less masking) for a $d=12$ transformer for the text in Appendix A.6.
Refer to caption: extracted/5899480/images/interp_triangles_thresh2.png

## 8 Related Works

Transformer Improvements. Our work aims at improving the transformer architecture.
Since its introduction in Vaswani et al. (), 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., ), gated units instead of MLPs (Shazeer, ), removing the biases (Raffel et al., ), using RMSNorm instead of LayerNorm (Zhang & Sennrich, ), normalizing the Q and K projections (Dehghani et al., ), and multi-query and group-query attention (Shazeer, ; Ainslie et al., ).

Attention Modifications. Our work focuses on modifying the attention module.
Most of the existing 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., ; Ding et al., ), and linear attention approximations (Shen et al., ; Katharopoulos et al., ; Schlag et al., ). Some works focus on hardware-aware optimizations instead, such as FlashAttention (Dao et al., ) and Ring Attention (Liu et al., ).

Context Pruning. A part of our work (Section 4) consists of removing elements from the context buffer. Several works employ this mechanism in order to increase inference efficiency.
Among those are *compression* methods, that aim to learn a compressed representation for tokens in the context buffer, with or without auxiliary losses, in order to replace several tokens with their compressed form (Munkhdalai et al., ; Ren et al., ; Yun et al., ; Mu et al., ; Rae et al., ).
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., ; Beltagy et al., ; Zaheer et al., ).
More sophisticated variants employ heuristics to evict less useful tokens from the context buffer instead of just the earliest ones (Oren et al., ; Zhang et al., ; Liu et al., ; Berchansky et al., ; Ge et al., ).
Dynamic context pruning (DCP) (Anagnostidis et al., ) proposed a mechanism for fine tuning existing models to make inference more efficient.
They learn which tokens to prune from the context buffer via a mechanism which shares similarities with selective attention.
DCP is more involved (e.g. needs root solving for evaluating the $\alpha$-sigmoid), introduces new parameters, requires an auxiliary loss, and produces binary prune decisions.
Notably, DCP didn’t observe the perplexity gains we see from selective attention(^3^33 We only became aware of DCP after this work was concluded. It is interesting that both works share some important similarities, although developed from very different perspectives (see Section 10).).

Inference Efficiency. A part of our work (Section 4) consists of making inference from transformers more efficient. Many approaches aim to speed up inference from transformers, including distillation (Hinton et al., ), sparcification (Jaszczur et al., ), quantization (Hubara et al., ), architecture modification (So et al., ; Shazeer, ), and algorithmic optimization (Dao et al., ; Leviathan et al., ).

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, ).

## 9 Discussion

In this work we introduced *Selective Attention*, a simple parameter-free change to the standard attention 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 a negligible amount of compute, and provides consistent improvements, selective attention might be a good default for transformer decoders.

Limitations and future directions.
We applied selective attention to decoder-only transformers. It could be interesting to investigate its applicability to encoders as well; Reducing the size of the context as in Section 4 improves inference efficiency 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 4. 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 attention.
It is interesting to investigate how it could be applied in a fine-tuning step to existing models;
While we observed similar results with selective attention in several setups (Section 5), there are still important variants we did not test, notably transformers with multi-query (Shazeer, ) and grouped-query (Ainslie et al., ) attention, as well as models much larger than 1B parameters;
Finally, it would be interesting to implement selective attention in a GPU-aware way, similar to Flash Attention (Dao et al., ).

## 10 Improving Neural Architectures

In *The Art of Transformer Programming*, Leviathan () 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 a mechanism allowing to selectively mask items in the context buffer, similar to selective attention. They further hypothesized that such a mechanism will have similar positive effects on language modeling, which motivated our work.

Zhou et al. () proposed the *RASP-Generalization Conjecture*, that “Transformers tend to length generalize on a task if the task can be 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 using RASP are hard for transformers as well, and if we made those easier, by changing the transformer architecture (and respectively the capabilities of RASP) we could meaningfully improve transformers.
Similarly, when constructing transformer programs by hand, Leviathan () 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 AI systems.”

We are very strong advocates for this method, and believe that finding basic problems for which we cannot program a general solution by hand on a neural model, is an extremely fertile approach for producing further 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.

## Appendix A Appendix

### A.1 Transformers with Selective Attention Learn a General Solution to Variable Assignment

Transformers with selective attention reach close to 0 validation loss and 100% precision extremely fast when trained on Variable Assignment 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 Assignment 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 a validation loss of 0.002 (and 100% accuracy) after less than 1,000 training steps. The transformer without selective attention only achieves a validation log-perplexity 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 of distribution test set with the same 3 variables but only 2 possible values, the transformer with standard attention’s accuracy drops substantially to 70% (loss of 3.64). Meanwhile the transformer with selective attention maintains 100% accuracy (with a loss of 2.4e-8).

We observed similar results in other settings (e.g. 10 variables and 10 possible values). We also repeated the experiments with somewhat larger transformers ($d=8$) and observed similar results.

### A.2 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 on C4 for 524,288 steps, to similarly trained transformers where the selection function is implemented via a separate bilinear form (adding additional parameters and computation).
The transformers with standard selective attention (i.e. sharing the outputs of an existing attention head) achieve the same or slightly better log-perplexities at the end of training (see Table 2).

**Table 2: The log-perplexity on the validation set after 524,288 training steps for (1) standard attention, (2) selective attention with a separate bilinear form for the selection module (more parameters than the baseline), and (3) selective attention.**
|  | $d=8$ | $d=12$ | 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.3 Ablating the Constraints

We ablate the 3 constraints selective attention applies to $S$ (see Section 3.2).

Negative selection. While with selective attention a token can decide to reduce attention to another token by all future attention operations, allowing a token to *strengthen* another token’s contribution to 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, ), it is plausible that masking it is detrimental. When we allow selective attention to mask the <BOS> token, we observe neutral to slightly worse results compared to standard selective attention where the <BOS> token is forced to never be selected, see Table 3.

Self-masking. Since selective attention reuses an existing attention head as the selection function, motivated by the absorption observation (see Section 3.1), it seems plausible that a token should never mask itself.
Indeed, when we allow tokens to mask themselves (i.e. we stop zeroing out the diagonal of $S$) we observe worse results, see Table 4.

**Table 3: The log-perplexity on the validation set after 524,288 training steps for (1) selective attention 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 4: The log-perplexity on the validation set after 524,288 training steps for (1) selective attention 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 |

### A.4 Self-Impact

Table 5 compares the results of forbidding self-impact (i.e. not allowing a token to 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 a small but consistent improvement.

**Table 5: The average log-perplexity on the validation set of 3 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.516 |
| Selective attention | 2.923 | 2.831 | 2.750 | 2.691 | 2.639 | 2.511 |

### A.5 Perplexity-Efficiency Trade-Off

Figure 6 illustrates the trade-off between perplexity gains and efficiency gains when pruning as in Section 4. See Section 6 for details.

Figure: 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 or equal to that of the baseline without selective attention (the dotted lines). Selective attention transformers trained with the $\mathcal{L}_{mem}$ loss and $\epsilon=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.
Refer to caption: extracted/5899480/images/tradeoffs.png

### A.6 Example Details

The following text from the C4 validation set was used in Figures 5, 9, and 10:

the real problem with traditional dental veneers has little to do with how they function or their performance . who determines what a normal , aesthetically pleasing smile looks like is the real issue . america struggled for decades with defining ‘‘ image ’’ as 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 as the standard to which we should aspire . the effects of this insidious marketing made their way into breast implants and the definition of a 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 people with nice smiles are convinced only a ‘‘ perfect ’’ hollywood smile is acceptable . one particularly relevant example of this is highlighted in a june 2015 article entitled ‘‘ saving jane ’ s smile . ’’ gary nankin , dds discusses how he ‘‘ saved ’’ the smile of a patient who was not content with her first set of # porcelain veneers . 1 . endodontic referral for treatment of tooth number 15 , followed by a 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 . 5 . placement of dental implant by the periodontist followed by preparation of mandibular teeth and placement of permanent restorations . 6 . restore the now fully - healed and osseointegrated implant in the position of tooth number 30 . regarding a person ’ s smile , the strong link to self - esteem and self - worth make an imperfect set of teeth a 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 ? .

### A.7 Comparison with Local Attention

Table 6 compares the validation perplexity of $d=12$ transformers with various local attention patterns (all-local and alternating), to that of a standard transformer and to that of a 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 a 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 3 local attention layers followed by 1 global attention layer, in a repeated fashion.
For example, “Local-global 32” denotes a transformer with 3 local attention layers where tokens can only attend up to 32 tokens back, followed by a 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 a transformer with selective attention.

**Table 6: Validation log-perplexity of transformers with different local attention patterns.**
| Model Type | Validation Log-Perplexity |
| --- | --- |
| 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.8 Additional Figures for Context Pruning

Figure 7 illustrates the values of the $F$ matrix averaged across 1,000 examples of length at least 512 from the C4 dataset.

Figure 8 compares the values from Figure 7 to those obtained from a different training run (different random initialization and different data shuffle). While this isn’t always the case, we sometimes observe stable sparsity patterns like those in this example, hinting at some general properties of language modeling on C4.

Figure 9 illustrates the tokens that remain in the context buffer after pruning (as in Section 4) for the example text (see Appendix A.6), for a $d=12$ model with selective attention, trained with the $\mathcal{L}_{mem}$ loss (Equation 2, $\epsilon=0.1$), for a memory budget where the validation perplexity matches a transformer without selective attention. The per-layer memory budgets chosen by the pruning algorithm for this model are: [8, 48, 8, 8, 24, 8, 168, 16, 8, 64, 8, 8], leading to a memory saving factor of 16X.

Figure 10 shows which tokens are pruned for the example text from Appendix A.6.

Figure: Figure 7: Visualization of the $F$ matrix (greener is lower, i.e. less masking) for a $d=12$ transformer averaged across 1,000 examples.
Refer to caption: extracted/5899480/images/interp_smooth.png

Figure: Figure 8: Visualization of the $F$ matrix (greener/bluer is lower, i.e. less masking) for a $d=12$ transformer averaged across 1,000 examples for two training runs (different random initialization, and different shuffle of the training data). While we only have anecdotal evidence, it is interesting that we sometimes observe these stable sparsity patterns across training runs.
Refer to caption: x2.png

Figure: Figure 9: Visualization of the persisted elements for a $d=12$ transformer for the text in Appendix A.6. The white pixels denote tokens removed from the context buffer as per Section 4.
Refer to caption: extracted/5899480/images/interp_triangles_2.png

Figure: Figure 10: A visualizing of the elements that are masked for the last (512th) token, for a $d=12$ transformer for the text in Appendix A.6. We observe some interesting patterns, for example, layer 4 persists the end-of-sentence periods (“.”).
Refer to caption: extracted/5899480/images/ashkara_detailed_2.png

### A.9 Parameter Counts

Table 7 shows the actual parameter counts for models with different $d$s (see Section 5).
Note that selective attention does not add any extra parameters.

**Table 7: The number of model parameters for different values of $d$ as in Section 5.**
| $d$ | $n_{layers}$ | $n_{heads}$ | $d_{model}$ | Number of parameters |
| --- | --- | --- | --- | --- |
| 8 | 8 | 8 | 512 | 33,603,584 |
| 10 | 10 | 10 | 640 | 59,699,200 |
| 12 | 12 | 12 | 768 | 97,615,872 |
| 14 | 14 | 14 | 896 | 149,666,048 |
| 16 | 16 | 16 | 1024 | 218,226,688 |
| 18 | 18 | 18 | 1152 | 305,717,760 |
| 20 | 20 | 20 | 1280 | 414,387,200 |
| 22 | 22 | 22 | 1408 | 546,641,920 |
| 24 | 24 | 24 | 1536 | 704,950,272 |
| 26 | 26 | 26 | 1664 | 891,468,032 |
| 28 | 28 | 28 | 1792 | 1,108,645,888 |

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
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?

2. 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?

3. 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?

4. 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?

5. 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?
soundness: 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.
2. **No Additional Parameters**: Selective Attention operates without introducing new parameters or significantly increasing computational overhead, which preserves the simplicity of the transformer architecture.
3. **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.
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.
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.
2. **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.
3. **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.
4. **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.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 1
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)?
soundness: 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.
- 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.
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.

The 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.
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.
- 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.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: None
soundness: 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
- The proposed method is easy to implement, adds no parameters, and saves memory overhead
- Interesting results and analysis
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.

Experiments 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).
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. 
- Besides a more comprehensive evaluation, I also encourage the authors to strengthen their findings by trying selective attention on larger-scale models and datasets

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 3
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?
soundness: 3
strengths: - The method is simple and intuitive
- Results are very strong and consistent, showing both accuracy improvements, and very impressive memory reduction
- Experiments are quite extensive
- 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*."
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.
weaknesses: While I like this paper a lot, I have several reservations.

First, 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.

Second, I am not sure I am fully convinced by some of the major claims in the paper.
 
(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.

(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. 

Third, 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.

Meta-review: ### Summary
This paper presents an elegant method for selective attention that prune tokens already absorbed in previous representations. Thus the pruned tokens can be evicted from the kv-cache to improve compute and memory usage.

### Strengths
The presented method is intuitive and simple. Results are very strong and consistent, showing both accuracy improvements and memory reduction. Visualizations of motivating examples are easy to understand.

### Weaknesses
The description of the method is not super clear, and it takes more than pass to fully understand it.

Decision: Accept (Poster)