SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ideal Generation Speed-up. In Fig. 7 we presented throughput in 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,
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.
Figure 15: Attention weight for different layers for the context switch example in Table 2. Color here indicates that the token can be attended to and does not correspond to the actual attention weight. Notice the casual masking and the three emerging dense triangular sub-matrices, especially in layers 7, 8, 9 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 a 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 et al. [2022]). Similarity, in this case, is measured by calculating the cosine similarity of the keys  $\mathbf{K}$  across different tokens. To test how feasible such a pruning strategy is in our scenario, we calculate for the keys  $\mathbf{K}$  of each token, the 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 a different prediction for each of the tokens in the input, and token similarity by itself is not good evidence for pruning. We finally highlight that even if the method from Bolya et al. [2022] achieved comparable results, computational benefits will be
Figure 16: Maximum potential speed-up from our method achieved by the homogeneity of the memory allocation.
Figure 17: Minimum cosine similarity to other non-pruned tokens in the sequence. We group the tokens, based on whether these were subsequently dropped or not. Results are averaged across samples and layers.
negligible, as no memory benefits can be achieved in this case and generation is performed one token at a 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 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. To better check if some tasks are affected more by sparsity compared to others, we provide per-task zeroshot performance, see Figure 19. We
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 a 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
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 a challenge dataset for evaluating commonsense natural language inference, piqa 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 a benchmark for commonsense reasoning. Complementary, we also present results for multirc which is a dataset of short paragraphs and multi-sentence questions, openbookqa that is a multiple-choice elementary-level science questions, sciq 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 a conversation. We believe some of these datasets better capture longer contexts and longer-range interactions.
MODEL
BATCH-SIZE	BASELINSE	50% SPARSITY	70% SPARSITY	90% SPARSITY
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 a 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' in Figure 20. 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 future, we 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 a particular language model, as it can be applied to any autoregressive transformer architecture.
Figure 20: Perplexity (lower is better) for different levels of sparsity. Perplexity is averaged across tokens with context size varying from 1 to 2048. Three subplots correspond to pruning three different base pythonia models.
D Discussion
We presented a simple approach that can be applied to any decoder-based Transformer autoregressive model. In an era where big foundation models are taking the community by storm, we presented a relatively inexpensive technique that can be applied to 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 a 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. We refer the reader to Section 4 and Appendix A for a complete description of the training protocol. We have also released the code as 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 on autoregressive generation models based on Transformers. Although 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.