Dynamic Context Pruning for Efficient and Interpretable Autoregressive Transformers
Authors: Sotiris Anagnostidis
Dario Pavllo
Luca Biggio
Lorenzo Noci
Aurelien Lucchi
Thomas Hofmann
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Transformers
Context-pruning
Efficient Transformer

Abstract: Autoregressive Transformers adopted in Large Language Models (LLMs) are hard to scale to long sequences. Despite several works trying to reduce their computational cost, most of LLMs still adopt attention layers between all pairs of tokens in the sequence, thus incurring a quadratic cost. In this study, we present a novel approach that dynamically prunes contextual information while preserving the model's expressiveness, resulting in reduced memory and computational requirements during inference. Our method employs a learnable mechanism that determines which uninformative tokens can be dropped from the context at any point across the generation process. By doing so, our approach not only addresses performance concerns but also enhances interpretability, providing valuable insight into the model's decision-making process. Our technique can be applied to existing pre-trained models through a straightforward fine-tuning process, and the pruning strength can be specified by a sparsity parameter. Notably, our empirical findings demonstrate that we can effectively prune up to 80\% of the context without significant performance degradation on downstream tasks, offering a valuable tool for mitigating inference costs. Our reference implementation achieves up to $2\times$ increase in inference throughput and even greater memory savings.

Dynamic Context Pruning for Efficient and Interpretable Autoregressive Transformers
Sotiris Anagnostidis
Dario Pavllo
Luca Biggio
Lorenzo Noci"
Aurelien Lucchi
Thomas Hofmann
"ETH Zürich
"ML, CSEM SA
"University of Basel
Abstract
Autoregressive Transformers adopted in Large Language Models (LLMs) are hard scale to long sequences. Despite several works trying to reduce their com- putational cost, most of LLMs still adopt attention layers between all pairs tokens in the sequence, thus incurring quadratic cost. In this study, present a novel approach that dynamically prunes contextual information while preserv- ing the model's expressiveness, resulting in reduced memory and computational requirements during inference. Our method employs a learnable mechanism that determines which uninformative tokens be dropped from the context at any point across the generation process. By doing so, our approach not only addresses performance concerns but also enhances interpretability, providing valuable insight into the model's decision -aking process. Our technique can be applied to existing pre-trained models through straightforward tine-tuning process, and the pruning strength can be specified by sparsity parameter. Notably, our empirical findings demonstrate that we can effectively prune up to 80% of the context without sig- nificant performance degradation on downstream tasks, offering a valuable tool for mitigating inference costs. Our reference implementation achieves up 2x increase in inference throughput and even greater memory savings.
Introduction
The introduction of Transformers [Vaswani et al., 2017] in Large Language Models (LLMs) has profoundly influenced the landscape Natural Language Processing (NLP), due to their appealing scaling properties [Kaplan et al., 2020] and their ability to train efficiently modern hardware architectures designed for extensive parallel computing. As L.LMs grow larger and more complex, the challenges associated with training and deploying them become challenging the quest for processing increasingly longer sequences, as pure self-attention layers scale quadratically sequence length during train and inference.
prominent. Especially
To address this limitation, several efforts focus on efficient implementations the attention mecha- nism on dedicated hardware [Dao al., 2022, Touvron al., 2023], algorithmic procedures to directly tackle the quadratic complexity. The latter direction has led to numerous variants sacrificing the generality of the standard attention mechanism in favor of more efficient alternatives [Tay et al.,
"Equal contribution. Correspondence sotirios anagnostidis @nff.cthz.ch.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

2020, Kitaev et al., 2020, Choromanski et al., 2020b, Katharopoulos et al., 2020, Zaheer et al., 2020, Shi al., 2021, Lin et al., 2022, Zhu and Soricut, 2021. Dai et al., 2020], some of which are illustrated in Fig. 1. Specifically, a large number of these methods focus either on sparsifying the attention weights, reducing the size of the available context to each token, or compressing the number of tokens to reduce the size of the attention matrix.
These methods, however, are inherently static, in the sense that each token is either forced attend to fixed pre-specified context window, or the input context is compressed to fixed dimensionaliiyy regardless of the information content of the input sequence. Furthermore, a performance gap still exists with respect to pure self- attention in many applications, thus implying the existence of a non-trivial trade-off between the span of the attention context and the model's capabilities [Dao et al., 2022, Sun et al., 2021, Beltagy et al., 2020].
To address these challenges, and enhance inference efficiency, while staying faithful to pure self- attention, we pose the following question:
Can we dynamically prune past content based on the available context, while preserving as much as possible the expressivity of the model?
In response this question, we introduce novel method for context pruning in Transformer-based decoder architectures. Our approach adds a minimal amount of additional training parameters that enable individual tokens to dynamically remove portions of the input sequence in layer- wise fashion. Once part of the context is removed, it is disregarded for the remaining part of the autoregressive generation process, leading reduced memory usage and computational requirements during inference. To this end, we also design a dynamic data structure that implements efficient insertion/removal of tokens from the context while supporting batched inference. In contrast to traditional methods relying on local or sparse attention, which may not capture the nuances and dynamic nature of the data over long contexts, ours leverages contextual cues to dynamically determine the relevance of the available information through a learned mechanism. This is achieved by making use of a sparse sigmoid function [Peters et al., 2019, Martins et al., 2020]. As demonstrated by our experimental evaluations, this allows us to extract and utilize essential details in a more adaptive and accurate manner. The degree of pruning can be effectively controlled through a hyperparameter that effectively accounts for the sparsity level.
Our technique serves as a modular building block for existing pre-trained models and can be easily integrated through a minimal fine-tuning stage. For our study, we focus on GPT-2 models [Radford et al., 2019] as they are publicly available and widely benchmarked, but due to the uniformity of modern architectures, our approach can be straightforwardly extended to any autoregressive Transformer. Moreover, since our method is based on context pruning. it can be seamlessly combined with other approaches aimed at improving inference efficiency, such as quantization, weight pruning, approximate attention, or other hardware optimizations.
We find that up to 80% of the context can be successfully pruned, with minimal deterioration in terms of perplexity and zero-shot performance, while requiring significantly fewer resources during inference. We showcase how these improvements can lead to measurable practical gains, by providing an efficient implementation that reduces memory usage for caching during token generation. More specifically, for larger context sizes we get up to 50% wall-time latency reduction for each generation step. while still decoding with up to 2x larger batch sizes, leading thus to significant performance benefits. These findings highlight the potential of context pruning as a powerful technique to enhance the efficiency and interpretability of Transformers in NLP.
2 Related Work
Despite exhibiting human-level performance on number of challenging tasks, LLMs are resource- intensive and inefficient. While the human brain consumes roughly the amount of energy equivalent to dim light bulb, top-performing GPT models require multiple GPUs with ~80GB of memory each for inference [Strubell et al., 2019, Frantar and Alistarh, 2023a]. Several research efforts have been focusing on improving their efficiency and memory requirements from several different angles.
Weight Pruning and Quantization. Modern L.MSs have high memory and compute requirements for both training and testing. To address this limitation, number of research efforts [Kwon et al.,

---

Layer specific
Current Token
Attended Tokens
Masked Tokens
Attention
Local Attention
Sparse Attention
Adaptively Sparse Attention
Figure I: Visualization of the causal attention weights associated with standard, local, sparse causal attention, and our approach. Adaptively sparse attention (rightmost) prunes weights dynamically for each token, and it does not impose any restricting inductive biases on the final attention structure.
2022, Frantar et al., 2023, Frantar and Alistarh, 2023b] have resorted to the established practice of weight pruning (Hassibi et al., 1993] to efficiently compress the original model to more manageable size. Remarkably, a large percentage of the original weights can be safely removed, resulting in only marginal perplexity growth [Bahl et al., 1983]. An alternative approach to reduce the memory and compute is quantization [Dettmers et al., 2022, Yao et al., 2022, Xiao et al., 2022, Frantar et al., 2022], which reduces the precision of the model's 's numerical representation. Quantization schemes [Dettmers al., 2022] enable 8-bit matrix multiplication for both feed-forward and attention projection layers resulting in significantly improved memory allocation without incurring any performance degradation.
Efficient Transformers and context pruning. One primary constraint of Transformer-based models is their quadratic complexity with respect to the length of the input sequence. Extensive research explores alternatives that exhibit sub-quadratic scaling, resulting in three main strategies [Lin al., 2022]. The first replaces the attention mechanism with an alternative operation that features more favourable scaling with the imput sequence length [Peng et al., 2021, Katharopoulos et al., 2020, Choromanski et al., 2020a, Schlag et al., 2021]. While several recent methods this category show promise, none have emerged as definitive winner, and most tate-of-theeart language models still rely on the standard attention mechanism [Touvron et al., 2023, Chowdhery et al., 2022]. The second approach proposed compress the length of the input context, controlling the complexity of the attention operation but unavoidably sacrificing potentially relevant information from the original input [Lee al., 2019, Wang al., 2020, Jaegle al., 2021]. ]. The third approach involves pruning the attention matrix, preventing each token from attending to every other token within the context [Zaheer al., 2020, Martins al., 2020, Lee al., 2023]. This line of research is motivated by the theoretical finding highlighting that sparse Transformers retain the expressivity of their dense counterparts [Yun al., 2020]. Many methods in this category employ specially designed attention masks that aim to zero out as many entries as possible, often based on principles of locality, randomness, or combination of both. The main drawback of these methods is their mostly static nature, meaning that every token is compelled to attend to fixed context window and disregard the rest of the context regardless of its specific role within the input sequence. Our approach falls within this last category, and enables dynamic sparsification of the attention matrix for decoder models, without resorting to any potentially restricting inductive biases about its structure.
Implementation Speed-up Recently, hardware-optimizzed implementations [Dao et al., 2022, Touvron et al., 2023] have been proposed with the aim of optimizing computational resources during the training phase of Transformers [Hoffmann et al., 2022]. On the other hand, as recent breakthroughs have led to widespread adoption of these models [Ouyang et al., 2022, OpenAl, 2023, Köpf et al., 2023], performance during inference becomes more relevant by the day. In decoder- based autoregressive Transformers, the backbone architecture of most current state-of-the-arr LLMs, inference involves evaluating and generating tokens one by one, using cached previous activations to avoid redundant computations.
contrast training, the inference is memory bound [Shazeer, 2019, Ivanov et al., 2021, Pope et al., 2022]. Compute is under-utilized, especially when deploying larger models, as the time required to transfer model parameters and activations to hardware memory far exceeds the actual computational time. This is further exacerbated by recent trends to ever-increase the model size and enable longer context windows. a result, batch decoding, a promising direction for more efficient utilization of hardware resources, is impeded.
3

---

Adeptively Sparee Attention at layer!
Nerory bufter for logers
Stondd duke in Io born (O,, oucted undiwiiine (D)
mranalorrrr
Seep
leaming
model
GC1
12:2
1004
Sape
Fandlrrerr
loaming
mar
Seep
leaminn
model
10.07
10.4
wgb
leomingg
Figure 2: We illustrate the state of the memory buffer at the start of each iteration for our proposed approach. Dropped tokens are irrelevant for any subsequent generation step and their cached activations are erased. Since self-attention is a set operation, the buffer (keys/values) of the dropped tokens can be reused by subsequent tokens, ensuring that the data structure is as packed as possible. Here "X" denotes tokens that are currently being dropped. Red boxes correspond to tokens that are already dropped, for which Key- Sey-Value values are no longer being cached.
3 Methodology
... Nlocab) where n is the Background. We operate on sequences of text tokens T E {0, I,..,,.. length of the sequence and Reycab is the vocabulary size. Tokens are embedded into X Rnxd using an embedding layer. where d is the embedding dimension of the model. When necessary, we use the superscript é 21. I.} to denote the representations and weights at different layers. One layer of the Transformer-decodee architecture [Vaswani et al., 2017] is defined as
X = MHAALayerNorm X x)) + X
X FF (LayerNorm(Xi) + X,
(2)
where MHA stands for Multi-head self-attention defined
MHA(X) = Concatenate (X), head, (X).... head, Wo
where
(3)
head, X) = SA (Q,, K;, Vs).
Here Q: XWW. Q: K, = XWK, k., and V XWy, v. are the queries, keys and values and SA denotes the single-head self-attention. The weight matrices WQ.. WK., Wy, Rd linearly project the input embedding into the head dimension p. Finally, Wo E dddd is the output projection. The feed-forward part of the Transformer is defined as
FF(X) OHt(W W
Fs:
(5)
where drr is nonlinearity, and W F1 Wr, are linear layers with typical dimensions WA, We, Rdxx4-d and W F. E R4-dxd, A final projection layer W E Rdxmmaahh is used to project back to the vocabulary space and predict the next token from the representations XL, We are focusing on Pre-LN [Xiong et al., 2020] decoder-only architectures, meaning that attention is causally masked, ie. every input token attends to the first i tokens in the input sequence. Conceptually, acts by predicting these attention masks using learned mechanism layer-wise manner, with the introduction of additional constraints to make sure causality is preserved (i.e. if token is dropped, will remain dropped in the future). During inference, however, our method can efficiently be implemented by erasing tokens from the key-value cache commonly adopted in autoregressive attention models.
method
Background: key-value cache. In autoregressive Transformers, inference can be optimized by reusing pre-computed activations (keys and values) to accelerate the sequential generation of to- kens [Ott et al., 2019, Vaswani al., 2018, Wolf et al., 2020], bringing down the computational cost to generate single token to (nn) from (n²?) (where is the sentence length). Most existing sparse attention techniques ignore the specifies of this process and focus on sparsifying each attention operation separately. As non-attended tokens can still be attended to by subsequent tokens, memory benefits are limited. By contrast, our approach is compatible with this setting, allowing us design an efficient batched data structure where dropped tokens effectively removed from the computation.

---

3.1 Adaptively Sparse Attention
We allow the network selectively drop parts of the context that are no longer required. An illustration of our proposed method can be seen in Fig. 2. At each layer, we introduce the param- eters W W Km Raxr for dimension r E R. that calculate the interaction queries and keys Kim"
Que
Q Kin E IRaxr, as Q = W 2n and K = X'W We then calculate the interaction of token with token j at layer & as:
T In and In = a (LRINTKiI + .
if < k
k.j
j, = k,
(6)
0, k,
where (-) denotes the sparse sigmoid function introduced in Section 3.2 and B R a scalar parameter per layer, that controls the initial sparsity as seen in Fig. 3 (right). Indices in Q K Raxr refer to the rows of the matrices. We can then modify the self-attention
O (K)) + Iog(()) V
SA(Q), K; Vf) = softmax
For k we set I2 O. which leads to masking entries in the self-attention, corresponding to the regular causal masking. We also impose that a token cannot drop itself, thus I = 1. We want to preserve information regarding the current token as its predictions are particularly important in determining the next token for the regular language modeling task that we considering. Small values of n.j impose partial masking of the corresponding token in the attention, and complete
masking occurs when I = 0. The cumulative product over tokens + + k in Eq. (6) imposes that dropping token (when a (.) - 0) has an irreversible effect, as it will remain dropped for all subsequent tokens, and hence for the remaining of the generation process. The complexity of the pruning logic is Onn d.r r n² r), which is lower than the one of the self-attention operation for r <<d.
Our mechanism allows layers to act independently, meaning that different sparsity patterns are encountered across layers. We also experimented with tying the model's dropping decisions with depth by imposing that a token dropped at given layer cannot be attended to in subsequent layers. However, we observed worse results and hence did not pursue this further. This is perhaps expected, given numerous results and interpretability studies regarding sparsity patterns of attention heads at different layers [Ramsauer al., 2020, Hao et al., 2021].
3.2 Sparse Sigmoid
In Eq. (6), we use o--), as sigmoid-like function to let the network decide when and what to drop. We favour binary decisions, leading to interaction values of either 0 or Inspired by the o-entmax function introduced in Peters et al. (2019). Martins et al. [2020], we define the a-sigmoid (based on the entropies proposed by Tsallis [1988]) as:
o (x) a-sigmoid(z) argmaxpejo.n (p-x x Ha(p)),
(8)
where
-1 ala- (p (1 (1- if +
(v)
-p ppoog (1 p) log(1 - p), if 1.
By varying during the training, we can control the sparsity in the network, i.c. regulate the softness of the pruning mechanism. In practice, we start from small values of a 1 and increase it according to cosine scheduler, as shown in Fig. 3. Small values of allow meaningful gradient signals to pass through the dropping mechanism, which is crucial at the beginning of training. On the other hand, larger values of a lead to sparse results desired during inference. We thus increase a to values leading to very sparse solutions, as illustrated in Fig. 3. In practice, during inference, we replace a(.)) with the step function, that corresponds to - co. We also initialize the biases parameters 3 in (6) to positive value, ensuring that tokens at the beginning of training have prior towards not being dropped. This strategy also facilitates fine-tuning existing pretrained models, as module
5

---

13.0
1.0
2.
Spares Spro
2.6
0.5
3.4
2.0
02
Triinn Susss
300
200
500
100
400
500
2003
Contodt Langon
Figure 3: (Left) We use a cosine scheduler to set the values of during training. (Middle) For values of a mappings of the a-sigmoid saturate at =1/(a - 1). During inference, we replace the a-sigmoid with a step function, that corresponds the case -> (Right) Distribution of k. for
different values of 3 with respect to the distance between the tokens k - j. For this depiction, we assume random normally distributed vectors inputs and randomly initialized weights W W K...* according to 'He' initialization [He et al., 2015].
will initially default close to the identity function. The a-sigmoid along with the training schedule on a allows for good signal propagation properties for the gradients [Noci et al., 2022]. We also explored using regular sigmoid with varying temperature [Kim et al., 2022], leading suboptimal nonbinary predictions and instabilities during training. Training with our sparse sigmoid also directly eliminates the need have any auxiliary network [Lee et al., 2023].
3.3 Regularized Objective
We augment the regular language modeling objective with regularization that incentivizes the network f to drop parts of the sequence. We fine-tune pretrained models, with parameters 6, using the objective:
L(O, T) = Lim (0, T) + L Leporeity city (9, T),
(10)
where
Lim (0, T) CE(fe fe T), shift(T))
(11)
the regular cross-entropy loss for the language modeling task based on the original and shifted input tokens T, and
Lsparsity (0,T) T) Ln(n 1)
(12)
the sparsity loss, encouraging the model to prune the context. In total (Ln(n - 1))/2 entries of IN are learned, as indicated in Eq. (6). We choose 0 to enforce different levels of sparsity. In general, for a current position in the context, we define as sparsity, the percentage of the previous tokens dropped, i.e. (tokens < i dropped)/i.
4 Experiments
We fine-tune pretrained GPT-2 models 1, that support context size of up to 1024 tokens, on a subset of the English Wikipedia 20220301.en and English bookcorpus datasets. keep a separate test set where we report perplexity after training. All models shown, for fair comparison, were fine-tuned using the same lightweight training setup as described in Appendix A. When using our adaptive sparse attention,
cosine scheduler for the parameter displayed in Fig. and specify Qua W Kiy More ablations regarding optimization and variations
r = 64 for the dimensions of W of our dropping mechanism are provided Appendix B. Unless otherwise stated, results refer to GPT-2-small models. We use the term dense for the regular GPT-2 models, fine-tuned without any additional W Que W Ki, parameters.
'We use the pretrained models and tokenizers from https://hugginntten co/, for the GPT-2- medium, large, xI] models. Here 50257.

---

Baselines. We compare against the baselines presented in Fig. 1. Local Attention refers to a causal attention mask, where each token attends to the previous k tokens in the sequence, including itself. This can also be interpreted as restricting the receptive field of the model. Sparse Attention refers to the baselines from Child et al. [2019], Lin et al. [2022], where each token attends to the tokens satisfying (1) [a/k]
jj/k] and (2) the tokens A 1, 2 k 1 1,,..... [ k - (numbering starts from zero). We fine-tune these baselines using the same aforementioned fine- ane-tuning procedure, for different choices of k. leading to different levels of sparsity, depending on the current context size.
Data structure. Real-world deployment of our approach exhibits numerous challenges due to the nature of batched generation. In particular, we highlight differences in prompt length (initial prefix), different final lengths (termination criteria), and uneven dropping of tokens across different sentences. Maximum performance is achieved when the key-value cache represented as a contiguous block of memory, and any masking resulting from padding or removed tokens ("holes") will result in a decrease in efficiency. To this end, we devise an efficient batched data structure that allows for efficient insertion and deletion of tokens (leveraging the set nature of the self-attention operation), while (i) allowing the underlying storage be processed as contiguous block of memory and (ti) ensuring that the load factor of the data structure is high enough to guarantee performance speed-up. More details are provided the Appendix
4.1 Results
Perplexity vs sparsity. We first study how context-pruning changes for different levels of sparsity in Fig. 4. Depending on the current context size, our method allows for up to 80% of the context to be successfully pruned, i.e. removed, with no performance loss in terms of perplexity (-0.085 average gain in perplexity when context size is 1000 tokens for 80.35% of sparsity compared to the dense counterpart). Our method also adapts to the current context size, meaning network trained with specific sparsity regularization exhibits different levels of sparsity depending on the current context size. Compared the baselines, our method exhibits consistently lower perplexity results for the same level of sparsity.
Donne
26
Sparse Attent
Local Alterzin
20
24
Ours
20
18
33
10
0.4
02
3.0
0.3
Spardtyy
Spanthy
Scanthy
Scarthy
Figure 4: Perplexity (lower is better) for different levels of sparsity. (Left) Overall perplexity averaged across tokens with context size varying from 1 to 1024. The three plots on the right show perplexity for different context sizes.
Zero-Shot Performance. To test general model capabilities and complement perplexity evaluations, we provide results on several zero-shot tasks [Dettmers et al., 2022] in Fig. 5. Similar trends hold overall; our approach retains even outperforms the performance of the dense baseline, even for cases with high sparsity. These tasks involve scenarios where the model is required to perform without any specific training prior exposure to the target domain. The results obtained validate that the models' general capabilities can be retained, even under high levels of sparsity.
Computational Analysis. We analyse the gains in terms of FLOPs and required memory when generating new sequences due to caching in Fig. 6. Our dropping mechanism introduces additional computational costs for the calculation of Q Kf int and the logic behind dropping via Eq. (6). Due to the relatively small chosen parameter r, i.e. the output dimension of the interaction weights W W W k these are nevertheless minimal. Although the raw FLOPs benefit when using sparse models does not seem very significant, as aforementioned, inference is predominately memory-bound. The attention thus takes significant proportion of real-time inference [Dao et al., 2022]. On the contrary, dense matrix multiplications used for all linear projections are very efficient. Memory
7

---

0.50
0.48
0.40
0.46
0.38
Local Attention
GPT-2-small
0.36
Dense
Sparse Altertion
0.40
GPT larrge
GPT-2-medium
Ours
0.38
GPT-2-x
Sparsity
Sparsity
Figure 5: Mean zero-shot accuracy (higher is better) for the WinoGrande, HellaSwag, PIQA, and LAMBADA datasets. As the sparsity of all methods depends on the context size, average the expected sparsity based on the lengths of the prefixes in these datasets. (Left) GPT-2-small models and (right) all GPT-2 models.
benefits, on the other hand, are substantial, as the memory required for caching is a linear function with respect sparsity, with negative slope. Sparser solutions will thus additionally allow to generate more sequences in a batched fashion. This is particularly relevant for bigger models, also longer sequences, where batch decoding is major challenge [Shazeer, 2019].
FLOPS for Context-size 1000
Cached Memory Context-size 1000
0.35
keys values attendon
0.006
embtediing byer
logit layer
0.004
feed/frrard
give
0.002
attennion
drop tokens
9.00
6.000
Increacing sparsity
Figure 6: (Left) Distribution of FLOPs for models with different levels of sparsity. Here, embedding- layer refers to the embedding of the input sequence to the representation X logits-layer to the projections of the final representation XL according feed-forward components, summed across the different layers, the current representation queries, keys, values and the final output projection, attention to the actual softmax operation and drop-tokens to additional compute required for calculating Q K and performing dropping via Eq. (6). (Right) Memory requirements when caching previous activations (keys and values). When implementing dropping, interaction keys K have to be additionally cached. int
the vocabulary size, feed-forward to the
to the projection of
Throughput. We demonstrate how reduced context and reduced memory requirements can lead to significant real-world time throughput in Fig. Initially, our pruned networks slower in terms of latency for small context lengths, because of the additional cost associated with the logic behind pruning. Nevertheless, they quickly surpass the dense baseline that struggles as the context size increases. This verifies the fact that although raw FLOPs benefits look unsubstantial, in fact, this leads significant gains due the specific memory profile Transformers' inference. Crucially, our pruned networks can support a much bigger batch size, leading to significant throughput gains. More specifically, for long context sizes, our GPT-2-small model offers additional 98% margin in throughput for loss in perplexity of only 0.316, with respect to the dense counterpart. Similarly, our GPT-2-medium model can yield 189% additional throughput for only 0.084 loss perplexity for context size of 1000 tokens. In particular, the same model (for 1.0) provides a higher throughput than GPT-2-small model, while achieving 3. 3.769 lower perplexity. context windows become larger by the day in state-of-n are models, we expect these gains to become even more relevant.

---

Average throughput
Throughput generation step
GPT2laage 1241000915
GFT3 medum
350 100700005
1333 TRRTESS
8000
20
Takencs
6000
18
16
4000
2000
Dena
Throughput when context size 1000
Time per generation step
Deplldy wien conteet: Mee 1200
26
24
OPT2
OPP3 medium
Optinal butch alzaa
16
GP13
14
meedure
casa x.20
2004 thr
base
256
512
1024
2048
4096
8192
16384
200
400
600
BD0
1000
Totengus
Contect size
Figure 7: We measure throughput using the optimal batch size NVIDIA RTX A5000 GPU. (Left) Throughput in terms of tokens per second for different models and different levels of sparsity (top) averaged across tokens for context sizes from to 1024 and (bottom) when the context size 1000 tokens. (Right) Average (top) throughput for varying context size for the GPT--2mmdium model and average (bottom) time per generation step for varying context size. As models require significantly less memory, a larger batch size can be accommodated, where large portions of the throughput gains can be attributed to.
Freftx Generated
Genorptod tokenss
8
PASCE
Hurber rd
8
3
200
420
603
300
1200
8
8
XPN
100%
KUNCT
DCT
MAX
M.
COON
muN
PAI
XDU
200
420
602
200
1200
Layar
Figure 8: (Top) Example of pruned tokens for layer 5 for the GPT-2-small model fine-tuned with 0.3 during generation. Most pruning is triggered by punctuation. (Bottom-left) We calculate the probability of tokens to be kept in the context based on the part of speech (POS) of the words they correspond to. (Bottom-middle) Most dropping caused by tokens corresponding punctuation, but distinct layers behave differently. (Bottom-right Example of the number of tokens pruned by the tokens' position id, for layers of GPT-2-small.

---

Interpretability. Fig. 8 provides insights into the interpretability aspect of the model's decision- making process. It is observed that token removal predominantly occurs when encountering stop words (punctuation), which aligns with the intuition that local information within sentence becomes less relevant after its completion. Furthermore, it worth noting that layers at varying depths exhibit distinct behaviours, reinforcing our rationale for dissecting token removal decisions across depth. The variance sparsity distribution across different depths indicates the necessity conducting additional interpretability research to obtain valuable insights in the interactions of the tokens within the model. We provide more insights towards this direction the Appendix C.
Discussion
We proposed Adaptively Sparse Attention, novel approach dynamically prune the context in decoder-only Transformer architectures. Our results indicate that our technique performs favourably compared to competitive baselines in terms of the ratio between perplexity and sparsity of the attention weights. Remarkably our approach also significantly reduces the computational and memory requirements without affecting its final performance. We practically showcase these benefits achieving more than double the throughput at cases. Adaptively sparse attention comes with two additional practical advantages: first, it can be seamlessly integrated into existing pre-trained models via a cheap fine-tuning step; second, it represents an orthogonal contribution to the burgeoning research line aimed at increasing the level of efficiency of modern LLMs. As such, we envision its combination with existing techniques like weight pruning and quantization to be a promising avenue for future research.
References
Lalit R Bahl, Frederick Jelinek, and Robert L Mercer. A maximum likelihood approach to continuous speech recognition. IEEE transactions on pattern analysis and machine intelligence, (2):179-190, 1983.
Iz. Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint 05X5v2000.001550, 2020.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 7432-7439, 2020.
Daniel Bolya, Cheng- Yang Fu, Xiaoliang Dai, Peizhao Zhang, Christoph Feichtenhofer, and Judy Hoffman. Token merging: Your vit but faster. arXiv preprint arXiv:2210.09461,, 2022.
Rewon Child, Scott Gray, Alce Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprint 1904::/904./055000 2019.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, David Belanger, Lucy Colwell, and Adrian Weller. Masked language modeling for proteins via linearly scalable long-context transformers, 2020a.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 147944 2020b.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.00 22211.2111, 2022.
Zihang Dai, Guokun Lai, Yiming Yang, and Quoe Le. tranee-transformer: Filtering out sequential redundancy for efficient language processing, Advances in neural information processing systems, 33:4271-4282. 2020.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances Neural Information Processing Systems, 35:16344-16359 2022.
10

---

Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Llm. int8 O: 8-bit matrix multiplication for transformers scale. arXiv preprint arXiv:2208.07333 2022.
Elias Frantar and Dan Alistarh. Massive language models can be accurately pruned in one-shot. arXiv preprini arX1v:2301.00774, 2023a.
Elias Frantar and Dan Alistarh. Sparsegpt: Massive language models can be accurately pruned in one-shot, 2023b.
Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. Gptq: Accurate post-training
quantization for generative pre-trained transformers. arXiv preprint arXiv:2210. 17323, 2022.
Elias Frantar, Sidak Pal Singh, and Dan Alistarh. Optimal brain compression: A framework for accurate post-training quantization and pruning, 2023.
Yaru Hao, Li Dong, Furu Wei, and Ke Xu. Self-attention attribution: Interpreting information interactions inside transformer. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 12963-12971, 2021.
Babak Hassibi, David G. Stork, and Gregory Wolff. Optimal brain surgeon and general network pruning, IEEE International Conference Neural Networks, pages 293-299 vol.1, 1993.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision, pages 1026-1034, 2015.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendrieks, Johannes Welbl, Aidan Clark, al. Training compute-optimal large language models. arXiv preprint arXiv:2203. 15556, 2022.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhus Yuanzhi Li, Shean Wang. Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint (0Xiv:2/06.09685.. 2021.
Andrei Ivanov, Nikoli Dryden, Tal Ben-Nun, Shigang Li, and Torsten Hoefler. Data movement is all you need: A case study on optimizing transformers. Proceedings of Machine Learning and Systems, 3:711-732, 2021.
Andrew Jaegle, Felix Gimeno, Andrew Brock, Andrew Zisserman, Oriol Vinyals, and Joao Carreira. Perceiver: General perception with iterative attention, 2021.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001 108311, 2020.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pages 5156-5165. PMLR, 2020.
Schoon Kim, Sheng Shen, David Thorsley, Amir Gholami, Woosuk Kwon, Joseph Hassoun, and Kurt Keutzer. Learned token pruning for transformers. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 784-794, 2022.
Nikita Kitaev, tukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint 2:XX:20000.04451. 2020.
Andreas Köpf, Yannic Kilcher, Dimitri von Rütte, Sotiris Anagnostidis, Zhi-Rui Tam, Keith Stevens, Abdullah Barhoum, Nguyen Minh Duc, Oliver Stanley, Richárd Nagyfi, et al. Openassistant conversations-ddmmcrrtiiii large language model alignment. arXiv preprint 27X7022304.007327, 2023.
Woosuk Kwon, Sehoon Kim, Michael W. Mahoney, Joseph Hassoun, Kurt Keutzer, and Amir Gholami. A fast post-training pruning framework for transformers, 2022.

---

Heejun Lee. Minki Kang, Youngwan Lee, and Sung Ju Hwang. Sparse token transformer with attention back tracking. In The Eleventh International Conference on Learning Representations. 2023.
Juho Lee, Yoonho Lee, Jungtaek Kim, Adam R. Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: A framework for attention-basse permutation-invariann neural networks, 2019.
Tianyang Lin, Yuxin Wang, Xiangyang Liu, and Xipeng Qiu. A survey of transformers. Al Open, 2022.
André Martins, António Farinhas, Marcos Treviso, Vlad Niculae, Pedro Aguiar, and Mario Figueiredo. Sparse and continuous attention mechanisms. Advances in Neural Information Processing Systems, 33:20989-21000 2020.
Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi. Signal propagation in transformers: Theoretical perspectives and the role of rank collapse. arXiv preprint (XXfv:2206.03122 2022.
OpenAL. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. fairseq: A fast, extensible toolkit for sequence modeling. arXiv preprint arXiv:1904.01038, 2019.
Long Ouyang, Jeffrey Wu, Xu Jiang. Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 35: 27730-27744, 2022.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring broad discourse context. arXiv preprint arXiv:1606.066 060311060003,, 2016
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A. Smith, and Lingpeng Kong. Random feature attention, 2021.
Ben Peters, Vlad Niculae, and André FT Martins. Sparse peuuence-to-sequenee models. arXiv preprint 1rXiv:1905.05702, 2019.
Reiner Pope. Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Anselm Levskaya, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. Efficiently scaling transformer inference. arXiv preprint arXiv:2211.0511 2022.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Scidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlović, Geir Kjetil Sandve, al. Hopfield networks is all you need. arXiv preprint urXiv22008.02217 2020.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99-106.
2021.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers, 2021.
Noam Shazeer. Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv: 10t::/9/1.0220 2019.
Han Shi, Jiahui Gao, Xiaozhe Ren, Hang Xu, Xiaodan Liang, Zhenguo Li, and James Tin-Yau Kwok. Sparsebert: Rethinking the importance analysis in self-attention. International Conference on Machine Learning, pages 9547-9557. PMLR, 2021.
12

---

Emma Strubell, Ananya Ganesh, and Andrew McCallum. Energy and policy considerations for deep learning in nlp. arXiv preprint arXiv: 02243, 2019.
Simeng Sun, Kalpesh Krishna, Andrew Mattarella-Micke, and Mohit lyyer. Do long-range language models actually use long-range context? arXiv preprint arXiv:2109.09115. 2021.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey.(2020). arXiv preprint cs. cs.LG/2009.067222 2020.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint S.Kiv:2302.13977 2023.
Constantino Tsallis. Possible generalization of boltzmann- statistics. Journal of statistical physics, 014479-487, 1988.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Ashish Vaswani, Samy Bengio, Eugene Brevdo, Francois Chollet, Aidan N Gomez, Stephan Gouws, Llion Jones, Łukasz Kaiser, Nal Kalchbrenner, Niki Parmar, al. Tensor2tensor for neural machine translation. arXiv preprint arXiv: 1803.007 2018.
Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. 2020.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Transformers:
natural language processing. In Proceedings of the 2020 conference on empirical methods in natural language processing: system demonstrations, pages 38-45, 2020.
Guangxuan Xiao, Ji Lin, Mickael Seznec, Julien Demouth, and Song Han. Smoothquant: Accurate and efficient post-training quantization for large language models. arXiv preprint arXiv:2211.10433 2022.
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International Conference on Machine Learning, pages 10524-10533, PMLR, 2020.
Zhewei Yao, Reza Yazdani Aminabudi, Minjia Zhang, Xiaoxia Wu, Conglong Li, and Yuxiong He. Zeroquant: Efficient and affordable post-training quantization for large-scale transformers. Advances in Neural Information Processing Systems, 35:27168-277
2022.
Chulhee Yun, Yin- Wen Chang, Srinadh Bhojanapalli. Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. O (n) connections are expressive enough: Universal approximability of sparse transformers. Advances in Neural Information Processing Systems, 33:13783-13799 13794, 2020.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang. Li Yang, al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33:17283-17297, 2020.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can machine really finish your sentence? arXiv preprint 10X.v:9905.07830. 2019.
Zhenhai Zhu and Radu Soricut. H-transformer-I Id: Fast one-dimensional hierarchical attention for sequences. arXiv preprint arXiv:2107.11906, 2021.
13

---

A Experimental Setup
We use the pretrained GPT-2 models from huggingface. Parameters of the architecture for these models are provided in Table 1.
Name
Parameters
Number of layers
Number of heads
Model Dimension d
GPT--2small
124M
12
12
768
GPT-2-medium
350M
16
1024
GPT-2-large
774M
36
20
1280
GPT-2-x/
1558M
48
25
1600
Table I: Parameters of the architecture for the GPT-2 models.
Training. We fine-tune pretrained models subset of the English Wikipedia 20220301.en and English bookcorpus datasets, for total of 25000 steps with a batch size of 6. The datasets provided by huggingface at https: http::/://uggingaaaec wwtpccc...o/cttaeeeeeeeee /tddddaatttttttiiiiaaas / //huggingf ace. co ht/ddaaatts/booooooooo small and medium models and 5e for the large and x/ models with the Adam optimizer. We do not use any weight decay or any scheduler for the learning rate. For the self-attention operations we use flash-attention as provided by the scaled_ sacllddo__ppoooea cecoo_prouut_atteeeee _attentionnnnnt in pytorch-2.0*.
and https:
respectively. We use a learning rate of 1e for the
Evaluation. For the zero-shot accuracy experiments, we report accuracy on the WinoGrande (Sak- aguchi et al., 2021], HellaSwag [Zellers al., 2019], PIQA [Bisk al., 2020] and LAMBADA (Pa- perno et al., 2016] datasets, similar to Dettmers et al. [2022]. As samples in these datasets have different lengths, we report as sparsity the mean sparsity over the samples for which predictions were made.
Efficient Memory Allocation. As explained in Section 4, the computational cost of our method is greatly affected by the underlying data structure used for representing the key-value cache. Conceptually, the data structure should implement the following methods (all batched):
(): inserts a new token (K, V, and K in))
get (): returns the keys/values added so far as contiguous block of memory, as well as a binary mask used to represent padding and potential gaps due to removed tokens.
remove (): given a binary mask of the same shape as that returned by get O), removes the specified tokens from the data structure.
Ideally, the insertion and deletion operations should be as efficient as possible, while guaranteeing that the memory block returned by get() is as packed as possible (high load factor). Additionally, the data structure should support dynamic resizing as more tokens are inserted. A simple (yet inefficient) baseline consists implementing the above interface as a dynamic array (i.e. a pre- allocated buffer that is dynamically resized once full) where erased tokens are simply masked out. Such an implementation, while correct, does not result in any memory and computation savings. Instead, motivated by the intuition that self-attention is set operation meaning that tokens do not need to be stored in sequential order - we recycle the memory slots of erased tokens. We insert new tokens at the leftmost available position in the data structure (Fig. 2), and ensure that the load factor (the ratio between the length of the longest sentence n and the capacity the buffer) always greater than specified value. We choose a load factor of 0.9 and dynamically consolidate the data structure when the effective load factor falls below this threshold. We also mention that the asymptotic cost of these operations does not have significant impact on the final performance, still dominated by the (n) cost of the self-attention mechanism (for a single generated token). In our implementation, both push () and remove () have cost of (n), while get () has of (1) since it simply returns a view of the memory buffer. We also experimented with asymptotically faster implementations for push () and remove() () (at Olog n) using a priority queue), but found these to be slower in practice due to inefficient use of the GPU.
the overall cost
http://www..roooo/////////////ee wwt.s/abbeeeeeeeeeeeeeereeeeeeee nn. functional. seaetiooal.scaeed.. scaled_dot.
org/
prnelucc_atteetiio html
14

---

B Training Results and Ablations
Training Curves. We provide an example of training curve in Fig. 9. Initially, the loss of the model is high, as at initialization many of the tokens are dropped, despite the introduced 33 parameters. Higher values 3 can significantly mitigate this phenomenon. We noticed however that very high values of 3 lead to worse final solutions in terms of achieved complexity for a given level of sparsity, For our experiments we initialize S 2.0, vA E {1,2,
800
ontoxt-dae
Language Modicing Loss
600
Sparse sigiaee
400
300
Nongruned
100
20000 25000
20000 Stops
25000
Stops
Stops
Figure 9: Training curve for the GPT-2-small model trained with regularizer 7 0.3.
Setting a. As the value of a in the sparse-sigmoid functions rises, solutions become sparser, as indicated by the sparsity in Fig. 9 (right). During inference, want replace the sparse-sigmoid function with a simple step function, thus we should make the sparse sigmoid as similar as possible to the step function during training, to mitigate any inconsistencies caused by functional differences. In practice, we found no benefit from increasing a to values larger than 8.0. The speed by which increased also plays significant role. Increasing too quickly does not allow for the new interaction parameters to be optimized correctly, leading to suboptimal solutions. We found that a cosine scheduler for 25000 steps enough lead solutions of adequate sparsity. We present results when optimizing for different number of steps in Fig. 10.
contert-si2 1000
1000
Steps: 1000
Language Modeling Loss
10
800
Steps: 5000
Steps: 8000
Steps:
8
Steps: 25000
6
200
10000 20000 30000 40000 50000 Staps
10000 20000 30000 40000 50000 Steps
10000 20000 30000 40000 50600
Staps
Figure 10: Training curves when a different number of total steps is executed, using the same cosine scheduler for the values of a.
Setting r. For our experiments, we used = 64 for the embedding dimensions of the interaction weights. We experimented with different dimensions in Fig. 11. Larger dimensions lead to higher sparsity for the same perplexity but also require more memory to store the interaction keys. We found r = 64 be a good compromise and tradeoff between extra memory and performance for the same sparsity.
15

---

context. 1000
500
2.5
Language Mcdeling Loss
400
2.0
300
1.5
200
100
0.0
16 32
64
128
256
512
1024
5000
10000
15000
20000
25000
Dimension
Steps
Figure 11: Training curves when varying the dimension r. (Left) Final language modelling loss and (right) sparsity during training. For different dimensions r, similar perplexity is achieved. Larger dimensions generally allow for sparser solutions for the same perplexity.
Propagation of Pruning with Depth. We also experimented with tying the dropping decisions with depth. In this case we modify Eq. (6) as:
Mp Am=st1 Inn and IN (1Qni (I /
30) iif , <
j 0, Of j
(13)
Dropping a token in this case at a layer & directly enforces that the token is dropped for all subsequent layers. Such choice is inspired by similar pruning techniques in Transformer encoder models [Bolya et al., 2022] that usually, reduce the number tokens with depth. We found, however, that this choice led to some disadvantages.
Firstly, we determined that sparsity is not monotonic function of depth. Similar results have been shown before [Ramsauer et al., 2020]. Typically, the middle layers can be more sparse compared to the first and last layers. Secondly, for deeper models (GPT-2-large and GPT-2-x/), we found that the large number of elements over which the cumulative product is taken over in Eq. (13), led to difficulties in learning. This can be perhaps expected given that no such training objective was optimized during the pre-training phase of these models. Finally, propagating the pruning decisions across depth significantly complicates the challenges for an efficient implementation.
Freezing Initial Parameters. We experimented with training just the interaction weights W Qm. W Ku E _dxr BB E I, leaving the rest of the parameters of the network frozen. This led average increase in the perplexity of 9.285 for the same levels of sparsity, showing that modifications to the network's weights/logic are still necessary.
C Additional Results
Sparsity per Layer. In Fig. 12 present the average level of sparsity per layer. In Fig. 13 present similarly the number of tokens kept per layer for different sizes of the initial context.
Tokens that Cause Pruning. In Fig. 8 we presented which kind of tokens cause the pruning to take place. For the same example and settings, we present here exactly which tokens are dropped and when, in Fig. 14.
To better understand how well the dropping mechanism works, we create an artificial example, where we concatenate together text from three distinct, independent contexts. More specifically, we concatenate together the texts, seen in Table 2. Inspecting the attention matrices Fig. 15, reveals that the network has largely discovered the switches in the context and leamed to ignore most of the preceding text if that comes from different context.
Context Switch.
16

---

Laye:31
D41-
Contarrate
Figure 12: Sparsity per layer for different levels of regularization Here averaging predictions for different context sizes ranging from to 1024.
Reg. 0.007
Rog. 0.91
Reg. 0:03
Rog. 0.1
Contost size
0000
303
800
2003
400
200
23
20
30
Figure 13: Sparsity per layer for different levels of regularization 7. Different colors indicate different initial un-pruned context sizes.
Lebanon, officially the Republic of Lebanon or Lehanese Republic, is country Western Asia. It is located between Syria to the north and cast and Israel to the south, while Cyprus lies to its west across the Mediterranean Sea: its location at the crossroads of the Mediterranean Basin and the Arabian hinterland has contributed to its rich history and shaped cultural identity of religious diversity.
A motorcycle (motorbike, bike, trike (if three- three-wheeled)) is or three-wheelee steered by handlebar from a saddle-style seat. Motorcycle design varies greatly a range of different purposes: long-distance travel, commuting, cruising. sport (including racing). and off-road riding. Motorcycling is riding motorcycle and being involved in other related social activity such as joining motorcycle club and attending motorcycle rallies.
vehicle
Zorba's 's Dance is instrumental by Greek composer Mikis Theodorakis. The song featured for the dance, which has become known sirtaki, in the 1964 film Zorba the Greek, for which Theodorakis wrote the soundtrack, and became renowned around the world. now commonly played and danced to in Greek tavernas. The film's 's track has since been recorded as standalone song by many different musicians from around the world.
Table 2: Concatenated contexts used for the context switch example in Fig. 15.
Ideal Generation Speed-up. In Fig. we presented throughput a realistic scenario where 'holes' formed in the memory buffer in an uneven fashion between samples in the batch. We expect that as we employ our method for larger models and contexts, the maximum feasible batch size will be reduced, mitigating partly this unwanted phenomenon. To evaluate the maximum possible speed-up.
17

---

Figure 14: We illustrate during generation which tokens cause pruning. The same layer and model is used as in Fig. 8. Arrows indicate decisions to prune. Nodes correspond to tokens, as determined by the tokenizer.
800008
5000000
1001320
Ligwr
La/W
LAyo
Figure 15: Attention weight for different layers for the context switch example in Table 2. Color here indicates that the token can be attended and does correspond to the actual attention weight. Notice the casual masking and the three emerging dense triangular sub-matrices, especially in layers 8, and 10.
we generate samples using the same prefix and the same sampling strategy across samples. The holes generated in this case are the same across the batch. We demonstrate the throughput achieved in this case in Fig. 16. The benefits in this case are lot more clear.
Investigating Dropped Tokens. In Transformer encoder models, it has been shown that similar tokens can be successfully pruned in deeper layers (e.g. Bolya al. [2022]). Similarity, in this case, is measured by calculating the cosine similarity of the keys K across different tokens. To test how feasible such a pruning strategy is our scenario, we calculate for keys of each token, minimum cosine distance to other tokens in the sequence. We then group these distances based on whether the token was subsequently dropped or not, in Fig. 17. We can see that tokens that have other tokens very similar to them in the sequence are pruned. However, tokens with no other similar ones in the sequences can also be dropped.
We experimented with the framework from Bolya et al. [2022], which has exhibited very successful results for transformer encoder models out of the box, i.e. without any additional fine-tuning. We found, nonetheless, that the perplexity quickly diverged, i.e. increased, even for small levels of sparsity. This constitutes another indication that pruning decoder models requires additional effort. Compared to encoder models, decoders make different prediction for each of the tokens in the input, and token similarity by itself not good evidence for pruning. We finally highlight that even if the method from Bolya al. [2022] achieved comparable results, computational benefits will be
18

---

Average throughput
Throughput generation step
14000
20
10000
18
8000
16
3713umal
Jaag
4000
Dena
2000
Throughput when context size 1000
Time per generation step
beplldy wien contact: Mew 1300
26
24
CPF2
ort:2taree
byr:2-medumm
0Fti2.mal
6x
4x
3x10
Optimal butch alsaa
16
14
2F1
2x10
casa
2004 thr
base
256
532
1024
2048
4096
8152
16384
327 68
200
400
600 Contect size
BDC
1000
Toberrus
Figure 16: Maximum potential speed-up from our method achieved by the homogeneity of the memory allocation.
Dropped Not dropped
5
Density
1
0.0
0.2
0.4
0.6
0.8
1.0
Casine distance
Figure 17: Minimum cosine similarity to other non- conppruned tokens in the sequence. We group the tokens, based on whether these were subsequentially dropped or Results are averaged across samples and layers.
negligible, as no memory benefits can be achieved in this case and generation is performed one token time.
Attended Tokens across Layers. We provide additional visualization of the attended tokens in Fig. 18. Different layers exhibit different levels of sparsity and attend to different tokens.
How sparsity affects specific NLP tasks Evaluating GPT models pretrained by language mod- elling objective is in general challenging and different techniques have been developed based on the final intended use. This challenge motivated us to showcase zeroshot accuracy to downstream tasks, in addition to evaluating perplexity for upstream tasks. better check if some tasks are affected more by sparsity compared to others, we provide per-task zeroshot performance, see Figure 19. We
19

---

CIC TOTAON
W1C INUES
EMI
MEC INO
EIM BYMAON
DARIE TIN LOTEN
SIAAN EIST LIVER LECNII 205
SUCP 35 YOUE CONTRRNN
whserile eS TAES
hASS AO 9IAR FO' Cossumation
USE M5 dytss
Uppss
DACC TOREEONN
TOT TII
EARM
LOON UIITe
Bandening
BYC USERNMATE
part
CII
planee
SIZAD
EORR plant
FORR
LEROS LIMITOOS
LECOMATION
MOTANN IOE EMI
LIA SOW CAN
EWAN
UNII UNLEROMNS
LRANERS parT
BIE LOOKA
LIS LIVEER LEEMENS
EMTT
TRESS
LEADES
BODDME
FIR INSAARTTIN
ISASE IISH EMI BIN INNO SONEANS
TRIEK
ENT DIEMEE
INTAISA BIR part of 7O0 MARIE ENE
CT MITRANA
SEATER
DISM EI UILY DEENSS lakgg
root
DIAEE BISS
DONIS BID TOADET UN
LNI1 and BUT deyen phant FOR CMT LIN INNO TREA
SINCTCO EST STOMETG
LINS ENVERONAT
paIT of KURSIN
ZEMM MYRKON
BANTE LIRE
LIST DIVER DESNNN
"S INg
ETAERRR ERE DOKM LISTE
lorg BOM InsME
EXCAEE IIEH EMER ERI DI TIRE
and her:
burdacoon
m1 the NAN SSNDS
CLDE de
STENETO
INTS ESSEIMOTE
DENTRR LINE
TO KERSINN i DERIONS
E SYRRINN
BENTER
DITTA (IST LIVER
LECAANT EISENDE
TARN CIGAR NORO ITROONS LNN BONH ERTEEE TIMER EST STANEE EIS TEAR TOM HUSOON
IRTA DF
MIRM
TEVARS
LhE RONUS
LERGS LOMEECNS LICOHASIO
EWRRY EINEE ENNIC LNISH
the ERENCTIC
DYE EIEENOSSSS BRANERS LIS LINES OF BO
are TOCTE
UII LETEE
LEAMNDS
EAGH SIS DIOO TRAOORS TIE SCOR LISSEN TISAR EOU OENNER IT long DOS TSSOO
DE PIRDS
IEWEES
(URR) MIS LOSSS
ESGER TRRISSIA
TROM
TO'S MARIE
meekiral
I GOdd MEEO BEROONS
in LISSOI LEIEE
THE BIER LER long
Dou levoed
Pare
PUSS IM
more
types
STV VSSS trees
ACC BITTIN
EOS CUNN SYEM TAAOUR
fart HIME O' STRISE Io Terg BOM (aSeed DNTAES EIEE WI2' ERTS DY TIRE
and
EIARY DOAS TINE TNDE
plarta
BUCT
D3 THEQURADEE
vogetab EE futa aYC hata ArE
gIann
EN USE PEE
IEEEEST OF conmett
from [ST STATES to TIEK ION HUSNE pare lnge
WIT BT more ENTION typec of UO0 tree
EWARY LIKEH TINE INEE
VOS tes they
CLCP NOST CSNSEENNN JENT uegethh es JI NYd hATE J008 LIITAAN
for INMOMAMTTTOON
T90E WElIE or
COSTICIC
from chards lisr BoSH (over) TRETE with THE
diforr. Uypss
SIGRN WERE LIITN (TEM SONTY
Figure 18: We visualize attended tokens across layers. Color indicates the ability to attend to a token and not the actual attention weight.
also include zeroshot results for additional tasks, that require long-range dependencies. It becomes clear that some tasks are affected more than others, according to our intuition.
Inference Compute An additional benefit of our pruning strategy is that we can accommodate longer initial contexts for the same hardware. Table 3 showcases the maximum potential context windows for different batch sizes. Our strategy accommodates inference with much longer contexts, for a fixed device.
Additional Models As the original context grows, pruning the context makes an increasingly significant difference, as more opportunities for pruning exist. Generally, enlarging the context length of Transformers is fascinating area of research with novel techniques being currently proposed. In our experiments, we are limited by the maximum positional encodings of the base GPT models. To evaluate how our method performs in larger contexts, we additionally present preliminary results
20

---

1.000
9.36
1307
6.205
013
9064
6270
maltic
spertockss
0.90
333
0.00
03100
2.15
0.08
612
0.05
342
Figure 19: Zeroshot performance for individual tasks and different sparsity levels. Different tasks capture different capabilities of the model and may be affected differently by the varying enforced sparsity in the context. Hellaswag is challenge dataset for evaluating commonsense natural language inference, piga evaluates physical commonsense knowledge, lambada openai is a reasoning challenge that evaluates the capabilities of computational models for text understanding by means of a word prediction task, winogrande is benchmark for commonsense reasoning. Complementary. present results for multire which is a dataset of short paragraphs and multi-sentence questions, openbookqa that is multiple-choice elementary-level science questions, seiq evaluates science exam questions and coqa measures the ability of machines to understand a text passage and answer a series of interconnected questions that appear in conversation. We believe some of these datasets better capture longer contexts and longer-range interactions.
also
MODEL
BASELINSE
SPARSITY
70% SPARSITY
90% SPARSITY
BATCH-SIZE
16
1729
3312
5434
16221
32
865
1616
2653
7901
64
433
784
1274
3742
Table 3: Maximum possible sequence length (number of tokens) for different batch sizes and different levels of sparsity for the GPT-2-xl models, on single NVIDIA RTX A5000. X% Sparsity denotes our model that uses adaptively sparse attention.
on pythia language models. Pythia uses rotary position embedding, as many current state-of-the-art LLMs. We present results on 'Perplexity vs sparsity' Figure 20. These concrete results demonstrate that our pruning technique works out of the box for different positional encodings well. Longer contexts supported by our base model (for pythia models that 2048) additionally allow for higher levels of sparsity for the same performance. In the future, intend to also evaluate larger models, potentially with the use of LoRA [Hu et al., 2021]. We again highlight that our approach is not specific to particular language model, it can be applied to any autoregressive transformer architecture.
Pythia-7M
Pythiu-150M
Pyth
37
14
Figure 20: Perplexity (lower is better) for different levels of sparsity. Perplexity is averaged across tokens with context size varying from to 2048. Three subplots correspond to pruning three different base pythia models.
21

---

D Discussion
We presented simple approach that can be applied any decoder- Transformer autoregressive model. In an era where big foundation models are taking the community by storm, we presented relatively inexpensive technique that can be applied these models and can lead to significant gains in terms of computational resources required to perform inference. We truly believe the balance between model performance and resource efficiency is crucial for the widespread adoption of large-scale language models, and our approach offers promising avenue for achieving this goal. We hope that our technique offers more interpretable predictions and inspires future research.
Reproducibility. We have taken multiple steps to ensure the reproducibility of our experiments. refer the reader to Section 4 and Appendix A for a complete description of the training protocol. We have also released the code part of the supplementary material, including scripts on how to reproduce our results. Additionally, trained models will be released for further research.
Limitations. Our work focuses autoregressive generation models based on Transformers. Al- though we argue that the same technique can be applied out-of-the-box to any such model, we focused on text generation models and specifically the GPT-2 family. Due to the uniformity of recent architectures, we expect that these results generalize to other popular models.
22


Reviews:
[{"review_id": "PpzcRPJRfN", "reviewer": "Reviewer_65eF", "summary": "Given the trend in large language models, it is a pretty important problem to search to efficient architectures. In this direction is the line of work to make the attention component efficient by introducing sparsity in the attention block and allowing every token to attend only a subset of the previous tokens.\n\nThe paper presents usage of Adaptively Sparse Attention whereby\n1) The network learns to drop parts of the context no longer required.\n2) The tokens dropped at every layer are independent, and different set of tokens might be chosen to be persisted at every layer.\n3) The network is trined using sparse sigmoid functions that introduce sparsity during training itself in contrast to some works that have tried to introduce sparsity only during inference.\n\nExperiments show that this methodology can lead to pretty strong performance (with minimal perplexity loss), even with 80% sparsity.", "questions": "- Was there any analysis done on which set of NLP tasks get affected the most with such sparsity?\n- For inference on a given device, do we have numbers on the increase in the max sequence length that can be supported?\n- Do these changes impact the training time throughput in any way?\n", "limitations": "None discussed in paper, and nothing important that I can think of.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The work presented in the paper is pretty innovative and impactful in this direction as finding efficient transformer architectures is key to sustain and grow the research and production usage of these large language models.\n\nThe model successfully exploits the fact that the attention matrix in these models are pretty sparse and additionally encourages that with sparse sigmoid-like function to mask out some tokens.\n\nExperiments show pretty strong performance even with 80% sparsity in the network.", "weaknesses": "Given that the model is performing pretty well in terms of perplexity even with huge sparsity, it would be interesting to perform an analysis of whether there is a class of NLP tasks the is significantly affected (like knowledge intensive tasks?), or perhaps translation or similar tasks where the structure of the input sentence matters if the model is focusing to forget stopwords very early.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive feedback and for acknowledging the impact of efficient inference in autoregressive models. Here, we discuss the comments raised.\n\n> Was there any analysis done on which set of NLP tasks get affected the most with such sparsity?\n\nSuch an analysis is indeed interesting and would enhance our intuition. We have provided some preliminary evidence (see our *global response*) on task-specific performance and on zeroshot performance on some additional tasks that we consider interesting in the high-sparsity regime. Thank you for raising this point, we will update the paper with a concrete discussion on this matter.\n\n> For inference on a given device, do we have numbers on the increase in the max sequence length that can be supported?\n\nWe agree with the reviewer that increasing the maximum sequence length during inference for a given device is one of the most promising outcomes of pruning the context. In our case, the maximum context length is limited by the maximum supported sequence length of the pre-trained GPT model, see also our *global response*. Provided a base model that supports larger sequence lengths, benefits given fixed hardware, are substantial (see Table 1 of attached PDF in our *global response*). We also present results for *pythia* models that already support longer context windows (see Figure 2 of the attached PDF in our *global response*).\n\n> Do these changes impact the training time throughput in any way?\n\nSee also our *global response*. Training FLOPs are only marginally affected. Still, we only finetune pre-trained models, to avoid even the smallest increase in training throughput.\n\nWe thank the reviewer again for all the questions. We truly believe that addressing them helped improve our work significantly."}, {"review_id": "q7Ub8yAmoc", "reviewer": "Reviewer_LhSf", "summary": "The authors propose a modified dynamic masking operation to the traditional multi-headed attention in transformers in order to allow models to learn to drop tokens at specific layers during training. In order to facilitate learning, they use a sparse sigmoid that is annealed to interpolate from a traditional sigmoid up to a step function in the infinite limit. They show that this dynamic sparse attention achieves lower perplexity at the same sparseness levels compared to other modified attention mechanisms, as well as better zero-shot accuracy and throughput. The proposed model shows promise for improving the efficiency and overcoming the quadratic time complexity of traditional dense transformers.", "questions": "- The $\\alpha$-sigmoid operator is a bit confusing to me. How are argmax values calculated efficiently and how are gradients propagated through the argmax operator in equation 8?\n- Do the baseline GPT-2 models use FlashAttention in their implementation? If not, how does the throughput of the dynamically pruned model compare to one with FlashAttention?\n- How sensitive is the training to the correct annealing of the $\\alpha$ parameter?", "limitations": "Yes", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "Improving the efficiency of vanilla dense transformers is an important problem, and being able to learn the sparsification operation rather than setting a static prior is an interesting direction. The proposed method is simple and computationally efficient, and the manuscript is well written and clear. The results demonstrate marked improvements over other sparse attention mechanisms at similar sparsity levels as well as increased throughput.", "weaknesses": "The main concerns I have are with how well this adaptive sparsity can be used when training from scratch and on long range dependency benchmarks such as Long Range Arena [1]. Since the experiments are initialized from GPT-2, all dense information is present, the model simply has to learn which information it can safely ignore. However, when training from scratch this optimization problem becomes much more difficult since falling into a local optimum with respect to pruning early on can reduce the model’s effectiveness in the future as the $\\alpha$ annealing increases towards a step function. It would be interesting to see results (even on small models) on how much utilizing this mechanism during from scratch training affects perplexity, given the same compute budget. \n\n[1]  Long Range Arena: A Benchmark for Efficient Transformers. 2020. Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, Donald Metzler", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for recognizing the importance of efficient inference and the novelty of dynamically pruning the context as a successful alternative. In the following, we take the opportunity to address the comments made.\n\n> How well this adaptive sparsity can be used when training from scratch\n\nPlease see also our *global response*. We do not train from scratch to avoid extra computational costs. Training from scratch is a possibility, and local minima can be avoided by using a schedule for the sparsity parameter. However, our method is **designed** for finetuning, which is the most compelling option.\n\n> Long range dependency benchmarks\n\nWe agree with the reviewer that a close examination needs to be performed to determine whether long-range dependencies can be preserved. See our *global response* for some additional results on benchmarks that better capture some of these interactions. Long Range Arena represents a suitable and attractive benchmark for encoder models, where computational benefits from sparsity are minimal. This is due to the lack of the possibility for a cumulative effect in dropping, as the one we have introduced through Eq. (6). We could not find a suitable benchmark for autoregressive models. As longer contexts become more relevant, however, we are sure that such benchmarks will emerge, as hinted by concurrent work [1]. \n\n[1] ​​Mohtashami, Amirkeivan, and Martin Jaggi. \"Landmark Attention: Random-Access Infinite Context Length for Transformers.\" arXiv preprint arXiv:2305.16300 (2023).\n\n> How are gradients propagated through the argmax operator in equation 8?\n\nThe forward pass is calculated efficiently via a bisection approach, by iteratively narrowing down the interval that contains the exact solution. The backward pass can then be computed independently. For details on the backward pass, we refer to Section 3.4 of [2]. As this is essential to our work, we will update the text with a concrete pointer to the section of the cited paper.\n\n[2] Peters, Ben, Vlad Niculae, and André FT Martins. \"Sparse sequence-to-sequence models.\" arXiv preprint arXiv:1905.05702 (2019).\n\n> Do the baseline GPT-2 models use FlashAttention in their implementation?\n\nFor all attention operations, we use flash-attention as provided by the `scaled_dot_product_attention` in *pytorch-2.02* (L462-463 in Appendix). Flash attention becomes increasingly more important as sequence length increases and so removing flash attention will further highlight gain with respect to the baselines. Thank you for pointing this out, we will highlight this in the main text.\n\n> How sensitive is the training to the correct annealing of the parameter?\n\nWe provide some results in Appendix B and more specifically in Figure 10. In short, increasing $\\alpha$ rapidly does not allow for the new interaction parameters to be properly tuned, i.e. the achieved sparsity is worse for the same perplexity. Increasing it too slowly leads to solutions of equal quality, but requires more compute for the finetuning phase.\n\n\nWe thank the reviewer again for all the questions, we will update the paper to highlight the points raised."}, {"review_id": "PFD36e6TvQ", "reviewer": "Reviewer_uWWw", "summary": "The paper introduced a novel inference strategy for transformer models that focuses on inference efficiency. Instead of retaining all context tokens throughout the entire inference process, they gradually eliminate tokens as they move deeper into the layers. To determine which tokens to drop, they trained some small linear layers to predict the remaining relevant context. Their experiments revealed that approximately 80% of the tokens could be safely discarded without much adverse impact on downstream task performance or perplexity. As a result, this approach significantly reduces the computational resources required for inference when the context length exceeds 500 tokens.", "questions": "- The paper introduces the use of sparse sigmoid to gradually enforce sparsity on the context tokens. It is important to understand why sparse sigmoid techniques were chosen and if they offer any specific advantages. Could other sparsity techniques, such as l0-regularization with hard-concrete distributions (Louizos et al., 2017), or techniques used in movement pruning (Sahn et al., 2020), achieve similar results?\n- There seems to be a typo in Figure 5. The green curve should be labeled as \"Sparse Attention\".\n- Figure 2 appears to be a bit confusing without proper legends. While it is clear that \"X\" denotes dropped tokens, the meaning of the red blocks is unclear. Could you provide an explanation or add appropriate legends?", "limitations": "The authors mentioned the limitation of the working being exclusively tested on autoregressive language models, and specifically GPT2 model family. The paper would be stronger if the author can show positive results on stronger base models. ", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The paper demonstrates excellent writing with a clear flow of ideas.\n- The authors introduce a unique and innovative data structure that enables batch operations involving masked tokens.\n- The experimental results indicate that in scenarios where long context (> 500 tokens) is involved in the inference process, the models can safely drop up to 80% of the tokens without any impact on perplexity.", "weaknesses": "- **The choice of downstream tasks raises questions:** The downstream tasks evaluated in Figure 5 primarily involve small context sizes, and Figure 7 indicates that a smaller context leads to reduced throughput compared to the standard dense model. It would be more persuasive to demonstrate that task performance remains intact when longer contexts are required, while simultaneously achieving gains in inference efficiency.\n- **Insufficient experiments with stronger base models:** It appears that the proposed method inevitably leads to performance degradation for larger and more capable base models, as evident from both Figure 5 and 7. This is likely because stronger base models are better at utilizing contexts, and additional contextual information enhances performance. To strengthen the paper's argument, the authors should include more results using larger and stronger base models (>1.5B parameters), such as LLaMA, Pythia, or even OPT.\n- **Evaluating generation quality:** Though the authors perform evaluations on language modeling with perplexity, it does not necessarily align with generation quality. If would be helpful if the authors could further provide evidence that dropping context tokens does not affect generation quality.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive feedback and the interesting points raised. We are glad to hear that they found our method description clear and easy to follow. In what follows ahead, we would like to take the opportunity to address the primary concern.\n\n> Choice of downstream tasks and evaluation of generation quality\n\nWe thank the reviewer for raising this point. We have provided more task-specific performance results on our *global response*. Specific tasks selected, better capture long-range dependencies and interactions across sentences and paragraphs.\n\n> Stronger base models\n\nWe agree with the reviewer that generalization to stronger base models is not always straightforward. In our *global response*, we provide experimental evidence that our finetuning method works *without any adaptation* for other models, *pythia* in this case. We hope that the additional results serve as a convincing argument that the proposed method is applicable to a large class of models. \n\n> Why sparse sigmoid?\n\nWe experimented with a few different ways to enforce sparsity, including $L_0$-regularization and using a sigmoid with a tunable temperature parameter [Kim et al., 2022]. We found that these techniques were unsuitable for enforcing exact sparsity, while also leading to training instabilities. The sparse sigmoid on the other hand, inspired by Peters et al. [2019], leads to exact zeros for $\\alpha > 1$, while also leading to well-behaved gradients. \n\nLouizon et al. propose the hard concrete distribution, which is obtained by “stretching” a binary concrete distribution and then transforming its samples with a hard-sigmoid. They also rely on a tunable temperature parameter to control the “softness” of the distribution, similar to the techniques we experimented with, that we found unsuccessful. We do not apply our sparsity regularization in the weights directly but on the products $(\\mathbf{Q}\\_{\\text{int}}^{\\ell})\\_n^T (\\mathbf{K}\\_{\\text{int}}^{\\ell})\\_j$. This is where we believe the instability is emerging from, as gradients are back-propagated through the whole network.\n\nSahn et al. select $\\text{top}\\_k$ entries that are active in the forward pass and use a straight-through estimator to update weights even if they are pruned in the forward pass. This pruning strategy is less dynamic, as the number of $k$ active elements is chosen in advance. In our experiments, we found that models trained with the same regularization exhibited different sparsity levels for different contexts, i.e. the model “understood” that some contexts may require more information to be preserved compared to others.\n\nThank you for the relevant references, we will include them in the related work.\n\n> Typo in Figure 5.\n\nThank you, we will fix this.\n\n> Legend for Figure 2.\n\n“X” denotes tokens that are currently being dropped (their cached Key-Value values still exist up to this point). Red boxes corresponds to tokens that are already dropped, for which Key-Value values are no longer being cached. We will update the legend, to make things more clear.\n\nWe thank the reviewer again for the interesting points raised, we will update the paper to discuss these in detail."}, {"review_id": "jmmtjxNsb8", "reviewer": "Reviewer_jYJo", "summary": "This paper proposes an efficient and interpretable context-based dynamic pruning method. They use additional query / key layers to generate dynamic attention masks and sparsify self-attention maps. They also introduce a sparse sigmoid and regularization term to control the sparsity. In experiments, they demonstrate the lowest performance degradation compared to previous static attention pruning methods at the same level of sparsity. The results of throughput and speed analysis show that the proposed method can achieve additional inference efficiency with minimal performance loss. They analyze the distribution of remaining contexts with respect to part of speech and the depth of the layer. They also demonstrate that the proposed method dynamically prune attention based on contexts through context switch experiments.", "questions": "Do you have expectations on how the zero-shot performance difference between the pruning strategies will change as the number of model parameters increases?", "limitations": "As written in the Limitations section, scalability studies on larger language models (>7B) would provide further insights into the dynamic pruning method.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The proposed method is well-motivated and easy to follow.\n- The proposed method efficiently increases interpretability and sparsity with minimal performance drop compared to previous research.\n- They provide an extensive analysis of the relationship between context length, throughput, and speed in various parameter sizes.\n- Based on the proposed pruning method, they devise efficient batched data structures for the optimized computation.", "weaknesses": "- Quantitative comparison of throughput and speed with local/sparse attention would contribute to a comprehensive understanding.\n- Further qualitative study of interpretability (Fig. 8) varying the sparsity level (gamma) would provide additional intuitive observations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the feedback and the interesting questions. We are glad that the reviewer found the method easy to follow, and we hope our work sparks future work in the direction of efficient and more interpretable inference. In the following, we address comments made.\n\n> Quantitative comparison of throughput and speed\n\nLocal and Sparse attention have a fixed logic in determining which tokens to drop. These methods then, require fewer FLOPs for the same level of sparsity. For completeness, we present here preliminary results on the throughput for different levels of sparsity for a GPT-2-xl model on an NVIDIA RTX A5000.\n\n```\n+------------------------------------------------------+\n|Sparsity \\ Model|Local Attention|Sparse Attention|Ours|\n+----------------+---------------+----------------+----+\n|       0.4      |      640      |       641      | 570|\n+----------------+---------------+----------------+----+\n|       0.6      |      850      |       851      | 765|\n+----------------+---------------+----------------+----+\n|       0.8      |      1391     |      1390      |1190|\n+------------------------------------------------------+\n```\nNumbers indicate throughput in tokens/s, when the original unpruned context length is 1000 tokens. Although Local and Sparse Attention mechanisms allow for slightly higher throughput, these are counteracted by a higher drop in performance. This is captured both when measuring upstream perplexity (Figure 4), and especially when evaluating zeroshot performance (Figure 5).\n\n> Further qualitative study of interpretability.\n\nInterpreting attention weights is in general a challenging task. Our sparse attention is in that regard unique for two reasons. Firstly, only a subset of the context tokens are being attended to, so attention weights are concentrated only on the more important parts of the context. This allows us to study the significance of each token, depending on if it is being dropped or not. Secondly, by analyzing which tokens trigger pruning, we can see better understand at what time and which parts of the context become irrelevant. We have included more results in Appendix C, which we hope might help towards grasping a better intuition. Specifically, we present the sparsity patterns per layer in Figures 12 and 13. We also visualize the mechanism of dropping in Figure 14, including in an artificial scenario where we expect dropping to occur, by concatenating paragraphs of different content, in Figure 15. We also provide some additional results on the embedding nature of the dropped tokens.  \n\nHowever, we agree with the reviewer that more possibilities exist to study and interpret the nature of this attention. We hope that our approach offers concrete advantages towards such an interpretation. Some related references [1, 2, 3, 4].\n\nPreliminary and complementary to Figure 8 (bottom-left) we present the alive probability for different parts of speech elements and models with different sparsity patterns. As $\\gamma$ increases, the alive probability of different parts of speech decays differently.\n\n```\n+---------------------------------------+\n|gamma| NOUN |  DET |PROPN| VERB |  ADV |\n+-----+------+------+-----+------+------+\n| 0.0 | 0.967| 0.940|0.985| 0.970| 0.982|\n+-----+------+------+-----+------+------+\n| 0.0 | 0.919| 0.763|0.964| 0.935| 0.932|\n+-----+------+------+-----+------+------+\n| 0.0 | 0.850| 0.532|0.884| 0.822| 0.783|\n+-----+------+------+-----+------+------+\n| 0.1 | 0.672| 0.305|0.775| 0.648| 0.587|\n+-----+------+------+-----+------+------+\n| 0.3 | 0.508| 0.170|0.621| 0.439| 0.354|\n+-----+------+------+-----+------+------+\n| 1.0 | 0.309|0.0851|0.439| 0.253| 0.212|\n+-----+------+------+-----+------+------+\n| 3.0 | 0.159|0.0506|0.274| 0.127| 0.135|\n+-----+------+------+-----+------+------+\n| 10. |0.0803|0.0338|0.132|0.0455|0.0553|\n+---------------------------------------+\n```\n\n[1] Chefer, Hila, Shir Gur, and Lior Wolf. \"Transformer interpretability beyond attention visualization.\" Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2021.\n\n[2] Bolya, Daniel, et al. \"Token merging: Your vit but faster.\" arXiv preprint arXiv:2210.09461 (2022).\n\n[3] Rigotti, Mattia, et al. \"Attention-based interpretability with concept transformers.\" International Conference on Learning Representations. 2021.\n\n> How will zero-shot performance change as the number of model parameters increases?\n\nIn our conducted experiments involving models with varying numbers of parameters, a consistent trend has emerged. Specifically, we have consistently observed that models employing pruned contextual components exhibit the capability to retain the performance levels of their unpruned counterparts, even when subjected to notably high sparsity levels. \n\nThe continuous evolution and advancement of efficient fine-tuning approaches [4, 5] gives us confidence that the extensibility of these techniques to even larger models (>7B) is within reach, despite computational constraints. Looking ahead, we are particularly enthusiastic, as this avenue of investigation holds substantial promise, particularly given its potential for effective deployment within real-world systems.\n\n[4] Hu, Edward J., et al. \"Lora: Low-rank adaptation of large language models.\" arXiv preprint arXiv:2106.09685 (2021).\n\n[5] Dettmers, Tim, et al. \"Qlora: Efficient finetuning of quantized llms.\" arXiv preprint arXiv:2305.14314 (2023).\n\nWe thank the reviewer for the interesting points raised, which will be discussed in the main text."}]

Author rebuttal:
We would like to thank all reviewers for taking the time to review our paper and for the valuable feedback. Here we address common points raised and present more experiments that we believe further strengthen our findings.

**How sparsity affects specific NLP tasks (@65eF, @LhSf, @uWWw)**

Evaluating GPT models pretrained by language modelling objective is in general *challenging* and different techniques have been developed based on the final intended use. This challenge motivated us to showcase zeroshot accuracy to downstream tasks, in addition to evaluating perplexity for upstream tasks. The tasks we selected are commonly used, see e.g. Dettmers et al. [2022], Frantar et al. [2023b] or other popular benchmarks [gpt4all](https://gpt4all.io/index.html), [ilm-eval](https://tju01.github.io/ilm-eval/#?benchmark=lm-evaluation-harness). We agree with the reviewers, however, that some tasks may be affected more by sparsity compared to others. For that reason, we provide per-task zeroshot performance in the attached PDF, see Figure 1. We also include zeroshot results for additional tasks, that require long-range dependencies. It becomes clear that some tasks are affected more than others, according to our intuition.

**Impact on training time (@65eF, @LhSf)**

The dropping mechanism introduces a *marginal* computational cost in terms of FLOPs, as highlighted in Figure 6 (left). Furthermore, our approach can be applied as a post-processing step given an initial model trained in a standard fashion, decreasing the training cost even further. Training from scratch is also a possibility, given a suitable decay schedule for the sparsity parameter, but fine-tuning is more compelling as it can be applied to existing pretrained models. Also note that one could tune our hyperparameter $\gamma$ during training dynamically, to achieve a desired level of sparsity, e.g. by monitoring the running average of the sparsity during training and increasing/decreasing it in small steps (similar to ADA [1] in GANs). As, for evaluation purposes, we wanted to acquire multiple models with different levels of sparsity, we did not do that.

**Additional Models and Inference Compute (@65eF, @uWWw, @jYJo)**

As the original context grows, pruning the context makes an increasingly significant difference, as more opportunities for pruning exist. This is clearly demonstrated by Figure 7 (bottom right). Generally, enlarging the context length of Transformers is a fascinating area of research with novel techniques being currently proposed, e.g. [2, 3]. In our experiments, we are limited by the maximum positional encodings of the base GPT models. To evaluate how our method performs in larger contexts, we additionally present preliminary results on *pythia* language models. *Pythia* uses rotary position embedding [4], as many current state-of-the-art LLMs, e.g. LLaMa-(1 and 2).
We present results on “Perplexity vs sparsity” in the attached PDF, in Figure 2. These concrete results demonstrate that *our pruning technique works out of the box for different positional encodings* as well. Longer contexts supported by our base model (for pythia models that is 2048) additionally allow for higher levels of sparsity for the same performance, in the case of Figure 2 in the attached PDF, that is perplexity. In the future, we intend to also evaluate larger models, potentially with the use of LoRA [5]. We again highlight that our approach is not specific to a particular language model, as it can be applied to any autoregressive transformer architecture.

An additional benefit of our pruning strategy is that we can accommodate longer initial contexts for the same hardware, as raised by reviewer **@65eF**. Table 1 in the attached PDF showcases the maximum potential context windows for different batch sizes. Our strategy accommodates inference with much longer contexts, for a fixed device.


[1] Karras, Tero, et al. "Training generative adversarial networks with limited data." Advances in neural information processing systems 33 (2020): 12104-12114.

[2] Press, Ofir, Noah A. Smith, and Mike Lewis. "Train short, test long: Attention with linear biases enables input length extrapolation." arXiv preprint arXiv:2108.12409 (2021).

[3] Chen, Shouyuan, et al. "Extending context window of large language models via positional interpolation." arXiv preprint arXiv:2306.15595 (2023).

[4] Su, Jianlin, et al. "Roformer: Enhanced transformer with rotary position embedding." arXiv preprint arXiv:2104.09864 (2021).

[5] Hu, Edward J., et al. "Lora: Low-rank adaptation of large language models." arXiv preprint arXiv:2106.09685 (2021).


Meta-review: This paper presents an efficient and interpretable context-based dynamic pruning method for transformers, which allows models to learn to drop tokens at specific layers during training. The authors introduce a sparse sigmoid and regularization term to control the sparsity and use additional query/key layers to generate dynamic attention masks and sparsify self-attention maps. The experiments demonstrate that the proposed method achieves lower perplexity at the same sparseness levels compared to other modified attention mechanisms, as well as better zero-shot accuracy and throughput. The results also show that in scenarios where long context (> 500 tokens) is involved in the inference process, the models can safely drop up to 80% of the tokens with little impact on perplexity.

The reviewers found the paper to be well-written, clear, and innovative. The proposed method is well-motivated and easy to follow, and the authors provide an extensive analysis of the relationship between context length, throughput, and speed in various parameter sizes. They also devise efficient batched data structures for optimized computation based on the proposed pruning method.

However, some concerns were raised regarding the choice of downstream tasks, the performance of the method when training from scratch, and the evaluation of generation quality. It would be interesting to see results on long-range dependency benchmarks and stronger base models (>1.5B parameters), as well as further qualitative studies of interpretability varying the sparsity level. 

Overall, the paper is technically solid and has high impact on NLP. The reviewers recommend acceptance of the paper.
Decision: Accept (spotlight)