SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We fine-tune pretrained GPT-2 models  ${}^{1}$  ,that support a context size of up to 1024 tokens,on a subset of the English Wikipedia 20220301.en and English bookcorpus datasets. We keep a separate test set where we report perplexity after training. All models shown, for a fair comparison, were fine-tuned using the same lightweight training setup as described in Appendix A. When using our adaptive sparse attention, we use a cosine scheduler for the  $\alpha$  parameter as displayed in Fig. 3 and specify  $r = {64}$  for the dimensions of  ${\mathbf{W}}_{Q\text{ int }}^{\ell },{\mathbf{W}}_{{K}_{\text{int }}}^{\ell }$  . More ablations regarding optimization and variations of our dropping mechanism are provided in 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  ${\mathbf{W}}_{{Q}_{\text{int }},{\mathbf{W}}_{{K}_{\text{int }}}$  parameters.
<sup>1</sup>We use the pretrained models and tokenizers from https://huggingface.co/, for the GPT-2-{small, medium, large, x1} models. Here  $n_{\mathrm{vocab}} = 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  $i$  attends to the tokens satisfying (1)  $\lfloor i / k \rfloor = \lfloor j / k \rfloor$  and (2) the tokens  $k - 1, 2 \cdot k - 1, \ldots, \lfloor i / k \rfloor \cdot k - 1$  (numbering starts from zero). We fine-tune these baselines using the same aforementioned fine-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 is 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 to be processed as a contiguous block of memory and  $(ii)$  ensuring that the load factor of the data structure is high enough to guarantee a performance speed-up. More details are provided in the Appendix A.
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 a network trained with specific sparsity regularization exhibits different levels of sparsity depending on the current context size. Compared to the baselines, our method exhibits consistently lower perplexity results for the same level of sparsity.
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 or 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 or 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  $\mathbf{Q}_{\mathrm{int}}^{\ell}$ ,  $\mathbf{K}_{\mathrm{int}}^{\ell}$  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  $\mathbf{W}_{Q_{\mathrm{int}}}^{\ell}$ ,  $\mathbf{W}_{K_{\mathrm{int}}}^{\ell}$ , 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 a 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
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, we 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 to sparsity, with a negative slope. Sparser solutions will thus additionally allow us to generate more sequences in a batched fashion. This is particularly relevant for bigger models, also longer sequences, where batch decoding is a major challenge [Shazeer, 2019].
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  $\mathbf{X}^0$ , logits-layer to the projections of the final representation  $\mathbf{X}^L$  according to the vocabulary size, feed-forward to the feed-forward components, summed across the different layers, qkvo-caclulation to the projection of the current representation to queries, keys, values and the final output projection, attention to the actual softmax operation and drop-tokens to additional compute required for calculating  $\mathbf{Q}_{\mathrm{int}}^{\ell}$ ,  $\mathbf{K}_{\mathrm{int}}^{\ell}$  and performing dropping via Eq. (6). (Right) Memory requirements when caching previous activations (keys and values). When implementing dropping, interaction keys  $\mathbf{K}_{\mathrm{int}}^{\ell}$  have to be additionally cached.
Throughput. We demonstrate how reduced context and reduced memory requirements can lead to significant real-world time throughput in Fig. 7. Initially, our pruned networks are 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 to significant gains due to the specific memory profile of 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 an additional  $98\%$  margin in throughput for a 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 in perplexity for a context size of 1000 tokens. In particular, the same model (for  $\gamma = 1.0$ ) provides a higher throughput than a GPT-2-small model, while achieving 3.769 lower perplexity. As context windows become larger by the day in state-of-the-art models, we expect these gains to become even more relevant.
Figure 7: We measure throughput using the optimal batch size on an 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 1 to 1024 and (bottom) when the context size is 1000 tokens. (Right) Average (top) throughput for varying context size for the GPT-2-medium model and average (bottom) time per generation step for varying context size. As our models require significantly less memory, a larger batch size can be accommodated, where large portions of the throughput gains can be attributed to.
Figure 8: (Top) Example of pruned tokens for layer 5 for the GPT-2-small model fine-tuned with  $\gamma -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 is caused by tokens corresponding to punctuation, but distinct layers behave differently. (Bottom-right) Example of the number of tokens pruned by the tokens' position id, for 2 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 a sentence becomes less relevant after its completion. Furthermore, it is worth noting that layers at varying depths exhibit distinct behaviours, reinforcing our rationale for dissecting token removal decisions across depth. The variance in sparsity distribution across different depths indicates the necessity of conducting additional interpretability research to obtain valuable insights in the interactions of the tokens within the model. We provide more insights towards this direction in the Appendix C.
We proposed Adaptively Sparse Attention, a novel approach to 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.
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 arXiv:2004.05150, 2020.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et 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, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 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.14794, 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.02311, 2022.
Zihang Dai, Guokun Lai, Yiming Yang, and Quoc Le. Funnel-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 in Neural Information Processing Systems, 35:16344-16359, 2022.
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-2-small	124M	12	12	768
GPT-2-medium	350M	24	16	1024
GPT-2-large	774M	36	20	1280
GPT-2-xl	1558M	48	25	1600
Table 1: Parameters of the architecture for the GPT-2 models.
Training. We fine-tune pretrained models on a subset of the English Wikipedia 20220301.en and English bookcorpus datasets, for a total of 25000 steps with a batch size of 6. The datasets are provided by huggingface at https://huggingface.co/datasets/wikipedia and https://huggingface.co/datasets/bookcorpus respectively. We use a learning rate of  $1e^{-4}$  for the small and medium models and  $5e^{-5}$  for the large and  $xl$  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_DOT_productattention in pytorch-2.0 $^2$ .
Evaluation. For the zero-shot accuracy experiments, we report accuracy on the WinoGrande [Sakaguchi et al., 2021], HellaSwag [Zellers et al., 2019], PIQA [Bisk et al., 2020] and LAMBADA [Paperno 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 we 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):
- push(): inserts a new token  $(\mathbf{K},\mathbf{V},$  and  $\mathbf{K}_{\mathrm{int}})$
- get(): returns the keys/values added so far as a 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(), 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 in 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 a 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 of the buffer) is always greater than a 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 a significant impact on the final performance, as the overall cost is still dominated by the  $\mathcal{O}(n)$  cost of the self-attention mechanism (for a single generated token). In our implementation, both push() and remove() have a cost of  $\mathcal{O}(n)$ , while get() has a cost of  $\mathcal{O}(1)$  since it simply returns a view of the memory buffer. We also experimented with asymptotically faster implementations for push() and remove() (at  $\mathcal{O}(\log n)$  using a priority queue), but found these to be slower in practice due to inefficient use of the GPU.
$^{2}$ https://pytorch.org/docs/stable/generated/torch(nn.Functional.scaled.dot_product attention.html
B Training Results and Ablations
Training Curves. We provide an example of a 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  $\beta_{\ell}$  parameters. Higher values of  $\beta^{\ell}$  can significantly mitigate this phenomenon. We noticed however that very high values of  $\beta^{\ell}$  lead to worse final solutions in terms of achieved complexity for a given level of sparsity. For our experiments we initialize  $\beta^{\ell} = 2.0, \forall \ell \in \{1, 2, \dots, L\}$ .
Figure 9: Training curve for the GPT-2-small model trained with a regularizer  $\gamma = 0.3$
Setting  $\alpha$ . As the value of  $\alpha$  in the sparse-sigmoid functions rises, solutions become sparser, as indicated by the sparsity in Fig. 9 (right). During inference, we want to 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  $\alpha$  to values larger than 8.0. The speed by which  $\alpha$  is increased also plays a significant role. Increasing  $\alpha$  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 was enough to lead to solutions of adequate sparsity. We present results when optimizing for a different number of steps in Fig. 10.
Figure 10: Training curves when a different number of total steps is executed, using the same cosine scheduler for the values of  $\alpha$ .
Setting  $r$ . For our experiments, we used  $r = 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$  to be a good compromise and tradeoff between extra memory and performance for the same sparsity.
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  $r$  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:
$$
\mathbf {I} _ {k, j} ^ {\ell} = \left\{ \begin{array}{l} \prod_ {p = 1} ^ {\ell} \prod_ {n = j + 1} ^ {k} \overline {{\mathbf {I}}} _ {n, j} ^ {p} \text {a n d} \overline {{\mathbf {I}}} _ {n, j} ^ {p} = \sigma \left(\frac {\left(\mathbf {Q} _ {\mathrm {i n t}} ^ {p}\right) _ {n} ^ {\top} \left(\mathbf {K} _ {\mathrm {i n t}} ^ {\ell}\right) _ {j}}{\sqrt {r}} + \beta^ {p}\right), \text {i f} j <   k \\ 1, \text {i f} j = k, \\ 0, \text {i f} j > k. \end{array} \right. \tag {13}
$$
Dropping a token in this case at a layer  $\ell$  directly enforces that the token is dropped for all subsequent layers. Such a choice is inspired by similar pruning techniques in Transformer encoder models [Bolya et al., 2022] that usually, reduce the number of tokens with depth. We found, however, that this choice led to some disadvantages.
Firstly, we determined that sparsity is not a 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-xl), 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  $\mathbf{W}_{Q_{\mathrm{int}}}^{\ell}, \mathbf{W}_{K_{\mathrm{int}}}^{\ell} \in \mathbb{R}^{d \times r}, \beta_{\ell} \in \mathbb{R}$ , leaving the rest of the parameters of the network frozen. This led to an 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.
Sparsity per Layer. In Fig. 12 we present the average level of sparsity per layer. In Fig. 13 we 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.
Context Switch. 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, as seen in Table 2. Inspecting the attention matrices in Fig. 15, reveals that the network has largely discovered the switches in the context and learned to ignore most of the preceding text if that comes from a different context.
Figure 12: Sparsity per layer for different levels of regularization  $\gamma$ . Here we are averaging predictions for different context sizes ranging from 1 to 1024.
Figure 13: Sparsity per layer for different levels of regularization  $\gamma$ . Different colors indicate different initial un-pruned context sizes.
Lebanon, officially the Republic of Lebanon or the Lebanese Republic, is a country in Western Asia. It is located between Syria to the north and east 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 a cultural identity of religious diversity.
A motorcycle (motorbike, bike, or trike (if three-wheeled)) is a two or three-wheeled motor vehicle steered by a handlebar from a saddle-style seat. Motorcycle design varies greatly to suit a range of different purposes: long-distance travel, commuting, cruising, sport (including racing), and off-road riding. Motorcycling is riding a motorcycle and being involved in other related social activity such as joining a motorcycle club and attending motorcycle rallies.
Zorba's Dance is an instrumental by Greek composer Mikis Theodorakis. The song featured for the dance, which has become known as sirtaki, in the 1964 film Zorba the Greek, for which Theodorakis wrote the soundtrack, and became renowned around the world. It is now commonly played and danced to in Greek tavernas. The film's track has since been recorded as a standalone song by many different musicians from around the world.
Table 2: Concatenated contexts used for the context switch example in Fig. 15.