SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sotiris Anagnostidis\*\*
Dario Pavllo\*\*
Luca Biggio $^{\mu, \nu}$
Lorenzo Noci
Aurelien Lucchi Thomas Hofmann
$\mu$  ETH Zurich
$\nu$ ML, CSEM SA
$^{\tau}$ University of Basel
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.
The introduction of Transformers [Vaswani et al., 2017] in Large Language Models (LLMs) has profoundly influenced the landscape of Natural Language Processing (NLP), due to their appealing scaling properties [Kaplan et al., 2020] and their ability to train efficiently on modern hardware architectures designed for extensive parallel computing. As LLMs grow larger and more complex, the challenges associated with training and deploying them become more prominent. Especially challenging is the quest for processing increasingly longer sequences, as pure self-attention layers scale quadratically in sequence length during train and inference.
To address this limitation, several efforts focus on efficient implementations of the attention mechanism on dedicated hardware [Dao et al., 2022, Touvron et al., 2023], or on 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@inf.ethz.ch.
2020, Kitaev et al., 2020, Choromanski et al., 2020b, Katharopoulos et al., 2020, Zaheer et al., 2020, Shi et 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 to attend to a fixed pre-specified context window, or the input context is compressed to a fixed dimensionality, 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 to this question, we introduce a 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 a layer-wise fashion. Once part of the context is removed, it is disregarded for the remaining part of the autoregressive generation process, leading to 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  $2\times$  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.
Despite exhibiting human-level performance on a number of challenging tasks, LLMs are resource-intensive and inefficient. While the human brain consumes roughly the amount of energy equivalent to a dim light bulb, top-performing GPT models require multiple GPUs with  $\sim 80\mathrm{GB}$  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 LLMs have high memory and compute requirements for both training and testing. To address this limitation, a number of research efforts [Kwon et al.,
Figure 1: 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 a 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 numerical representation. Quantization schemes [Dettmers et 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 et al., 2022]. The first replaces the attention mechanism with an alternative operation that features more favourable scaling with the input sequence length [Peng et al., 2021, Katharopoulos et al., 2020, Choromanski et al., 2020a, Schlag et al., 2021]. While several recent methods in this category show promise, none have emerged as a definitive winner, and most state-of-the-art language models still rely on the standard attention mechanism [Touvron et al., 2023, Chowdhery et al., 2022]. The second approach proposed to 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 et al., 2019, Wang et al., 2020, Jaegle et al., 2021]. The third approach involves pruning the attention matrix, preventing each token from attending to every other token within the context [Zaheer et al., 2020, Martins et al., 2020, Lee et al., 2023]. This line of research is motivated by the theoretical finding highlighting that sparse Transformers retain the expressivity of their dense counterparts [Yun et 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 a combination of both. The main drawback of these methods is their mostly static nature, meaning that every token is compelled to attend to a 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-optimized 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, OpenAI, 2023, Kopf 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-art LLMs, inference involves evaluating and generating tokens one by one, using cached previous activations to avoid redundant computations. In contrast to 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. As a result, batch decoding, a promising direction for more efficient utilization of hardware resources, is impeded.
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 can be 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-Value values are no longer being cached.
Background. We operate on sequences of text tokens  $\mathbf{T} \in \{0,1,\dots,n_{\mathrm{vocab}}\}^n$ , where  $n$  is the length of the sequence and  $n_{\mathrm{vocab}}$  is the vocabulary size. Tokens are embedded into  $\mathbf{X}^0 \in \mathbb{R}^{n \times d}$  using an embedding layer, where  $d$  is the embedding dimension of the model. When necessary, we use the superscript  $\ell \in \{1,2,\dots,L\}$  to denote the representations and weights at different layers. One layer of the Transformer-decoder architecture [Vaswani et al., 2017] is defined as
$$
\mathbf {X} = \operatorname {M H A} \left(\operatorname {L a y e r N o r m} \left(\mathbf {X} ^ {\ell - 1}\right)\right) + \mathbf {X} ^ {\ell - 1}, \tag {1}
$$
$$
\mathbf {X} ^ {\ell} = \operatorname {F F} (\operatorname {L a y e r N o r m} (\mathbf {X})) + \mathbf {X}, \tag {2}
$$
where MHA stands for Multi-head self-attention defined as
$$
\operatorname {M H A} (\mathbf {X}) = \text {C o n c a t e n a t e} \left(\operatorname {h e a d} _ {1} (\mathbf {X}), \operatorname {h e a d} _ {2} (\mathbf {X}), \dots , \operatorname {h e a d} _ {h} (\mathbf {X})\right) \mathbf {W} _ {O}, \quad \text {w h e r e} \tag {3}
$$
$$
\operatorname {h e a d} _ {i} (\mathbf {X}) = \operatorname {S A} \left(\mathbf {Q} _ {i}, \mathbf {K} _ {i}, \mathbf {V} _ {i}\right). \tag {4}
$$
Here  $\mathbf{Q}_i = \mathbf{X}\mathbf{W}_{Q_i}$ ,  $\mathbf{K}_i = \mathbf{X}\mathbf{W}_{K_i}$ , and  $\mathbf{V} = \mathbf{X}\mathbf{W}_{V_i}$  are the queries, keys and values and SA denotes the single-head self-attention. The weight matrices  $\mathbf{W}_{Q_i}$ ,  $\mathbf{W}_{K_i}$ ,  $\mathbf{W}_{V_i} \in \mathbb{R}^{d \times p}$  linearly project the input embedding into the head dimension  $p$ . Finally,  $\mathbf{W}_O \in \mathbb{R}^{d \times d}$  is the output projection. The feed-forward part of the Transformer is defined as
$$
\operatorname {F F} (\mathbf {X}) = \sigma_ {\mathrm {F F}} \left(\mathbf {X} \mathbf {W} _ {F _ {1}}\right) \mathbf {W} _ {F _ {2}}, \tag {5}
$$
where  $\sigma_{\mathrm{FF}}$  is a nonlinearity, and  $\mathbf{W}_{F_1}, \mathbf{W}_{F_2}$  are linear layers with typical dimensions  $\mathbf{W}_{F_1} \in \mathbb{R}^{d \times 4 \cdot d}$  and  $\mathbf{W}_{F_2} \in \mathbb{R}^{4 \cdot d \times d}$ . A final projection layer  $\mathbf{W}_{\mathrm{logits}} \in \mathbb{R}^{d \times n_{\mathrm{vocab}}}$  is used to project back to the vocabulary space and predict the next token from the representations  $\mathbf{X}^L$ . We are focusing on Pre-LN [Xiong et al., 2020] decoder-only architectures, meaning that attention is causally masked, i.e. every input token  $i$  attends to the first  $i$  tokens in the input sequence. Conceptually, our method acts by predicting these attention masks using a learned mechanism in a layer-wise manner, with the introduction of additional constraints to make sure causality is preserved (i.e. if a token is dropped, it 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.
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 tokens [Ott et al., 2019, Vaswani et al., 2018, Wolf et al., 2020], bringing down the computational cost to generate a single token to  $\mathcal{O}(n)$  from  $\mathcal{O}(n^2)$  (where  $n$  is the sentence length). Most existing sparse attention techniques ignore the specifics 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 to design an efficient batched data structure where dropped tokens are effectively removed from the computation.
3.1 Adaptively Sparse Attention
We allow the network to 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 parameters  $\mathbf{W}_{Q_{\mathrm{int}}}^{\ell}, \mathbf{W}_{K_{\mathrm{int}}}^{\ell} \in \mathbb{R}^{d \times r}$  for dimension  $r \in \mathbb{R}$ , that calculate the interaction queries and keys  $\mathbf{Q}_{\mathrm{int}}^{\ell}, \mathbf{K}_{\mathrm{int}}^{\ell} \in \mathbb{R}^{n \times r}$ , as  $\mathbf{Q}_{\mathrm{int}}^{\ell} = \mathbf{X}^{\ell} \mathbf{W}_{Q_{\mathrm{int}}}^{\ell}$  and  $\mathbf{K}_{\mathrm{int}}^{\ell} = \mathbf{X}^{\ell} \mathbf{W}_{K_{\mathrm{int}}}^{\ell}$ . We then calculate the interaction of token  $k$  with token  $j$  at layer  $\ell$  as:
$$
\mathbf {I} _ {k, j} ^ {\ell} = \left\{ \begin{array}{l} \prod_ {n = j + 1} ^ {k} \overline {{\mathbf {I}}} _ {n, j} ^ {\ell} \text {a n d} \overline {{\mathbf {I}}} _ {n, j} ^ {\ell} = \sigma \left(\frac {\left(\mathbf {Q} _ {\mathrm {i n t}} ^ {\ell}\right) _ {n} ^ {\top} \left(\mathbf {K} _ {\mathrm {i n t}} ^ {\ell}\right) _ {j}}{\sqrt {r}} + \beta^ {\ell}\right), \text {i f} j <   k \\ 1, \text {i f} j = k, \\ 0, \text {i f} j > k, \end{array} \right. \tag {6}
$$
where  $\sigma (\cdot)$  denotes the sparse sigmoid function introduced in Section 3.2 and  $\beta^{\ell}\in \mathbb{R}$  a scalar parameter per layer, that controls the initial sparsity as seen in Fig. 3 (right). Indices in  $\mathbf{Q}_{\mathrm{int}}^{\ell},\mathbf{K}_{\mathrm{int}}^{\ell}\in$ $\mathbb{R}^{n\times r}$  refer to the rows of the matrices. We can then modify the self-attention
$$
\operatorname {S A} \left(\mathbf {Q} _ {i} ^ {\ell}, \mathbf {K} _ {i} ^ {\ell}, \mathbf {V} _ {i} ^ {\ell}\right) = \operatorname {s o f t m a x} \left(\frac {\mathbf {Q} _ {i} ^ {\ell} \left(\mathbf {K} _ {i} ^ {\ell}\right) ^ {\top}}{\sqrt {p}} + \log \left(\mathbf {I} ^ {\ell}\right)\right) \mathbf {V} _ {i} ^ {\ell}. \tag {7}
$$
For  $j > k$  we set  $\mathbf{I}_{k,j}^{\ell} = 0$ , 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  $\mathbf{I}_{k,k}^{\ell} = 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 are considering. Small values of  $\overline{\mathbf{I}}_{n,j}^{\ell}$  impose partial masking of the corresponding token in the attention, and complete masking occurs when  $\overline{\mathbf{I}}_{n,j}^{\ell} = 0$ . The cumulative product over tokens  $j + 1 \rightarrow k$  in Eq. (6) imposes that dropping a token (when  $\sigma(.) \rightarrow 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  $\mathcal{O}(n \cdot d \cdot r + n^2 \cdot 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 a 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 et al., 2020, Hao et al., 2021].
3.2 Sparse Sigmoid
In Eq. (6), we use  $\sigma(\cdot)$ , as a 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 1. Inspired by the  $\alpha$ -entmax function introduced in Peters et al. [2019], Martins et al. [2020], we define the  $\alpha$ -sigmoid (based on the entropies proposed by Tsallis [1988]) as:
$$
\sigma (x) = \alpha \text {- s i g m o i d} (x) = \operatorname {a r g m a x} _ {p \in [ 0, 1 ]} \left(p \cdot x + H _ {\alpha} (p)\right), \tag {8}
$$
$$
H _ {\alpha} (p) = \left\{ \begin{array}{l} \frac {1}{\alpha (\alpha - 1)} (p - p ^ {\alpha} + (1 - p) - (1 - p) ^ {\alpha}), \text {i f} \alpha \neq 1 \\ - p \log p - (1 - p) \log (1 - p), \text {i f} \alpha = 1. \end{array} \right. \tag {9}
$$
By varying  $\alpha$  during the training, we can control the sparsity in the network, i.e. regulate the softness of the pruning mechanism. In practice, we start from small values of  $\alpha = 1$  and increase it according to a cosine scheduler, as shown in Fig. 3. Small values of  $\alpha$  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  $\alpha$  lead to sparse results desired during inference. We thus increase  $\alpha$  to values leading to very sparse solutions, as illustrated in Fig. 3. In practice, during inference, we replace  $\sigma(\cdot)$  with the step function, that corresponds to  $\alpha \to \infty$ . We also initialize the biases parameters  $\beta^{\ell}$  in (6) to a positive value, ensuring that tokens at the beginning of training have a prior towards not being dropped. This strategy also facilitates fine-tuning existing pretrained models, as our module
Figure 3: (Left) We use a cosine scheduler to set the values of  $\alpha$  during training. (Middle) For values of  $\alpha > 1$ , mappings of the  $\alpha$ -sigmoid saturate at  $\pm 1 / (\alpha - 1)$ . During inference, we replace the  $\alpha$ -sigmoid with a step function, that corresponds to the case  $\alpha \to \infty$ . (Right) Distribution of  $\mathbf{I}_{k,j}^{\ell}$  for different values of  $\beta^{\ell}$  with respect to the distance between the tokens  $k - j$ . For this depiction, we assume random normally distributed vectors as inputs and randomly initialized weights  $\mathbf{W}_{Q_{\mathrm{int}}}^{\ell}, \mathbf{W}_{K_{\mathrm{int}}}^{\ell}$ , according to 'He' initialization [He et al., 2015].
will initially default close to the identity function. The  $\alpha$ -sigmoid along with the training schedule on  $\alpha$  allows for good signal propagation properties for the gradients [Noci et al., 2022]. We also explored using a regular sigmoid with a varying temperature [Kim et al., 2022], leading to suboptimal nonbinary predictions and instabilities during training. Training with our sparse sigmoid also directly eliminates the need to have any auxiliary network [Lee et al., 2023].
3.3 Regularized Objective
We augment the regular language modeling objective with a regularization that incentivizes the network  $f$  to drop parts of the sequence. We fine-tune pretrained models, with parameters  $\theta$ , using the objective:
$$
L (\theta , \mathbf {T}) = L _ {l m} (\theta , \mathbf {T}) + L _ {\text {s p a r s i t y}} (\theta , \mathbf {T}), \tag {10}
$$
$$
L _ {l m} (\theta , \mathbf {T}) = \operatorname {C E} \left(f _ {\theta} (\mathbf {T}), \operatorname {s h i f t} (\mathbf {T})\right) \tag {11}
$$
is the regular cross-entropy loss for the language modeling task based on the original and shifted input tokens  $\mathbf{T}$ , and
$$
L _ {s p a r s i t y} (\theta , \mathbf {T}) = \gamma \frac {2}{L n (n - 1)} \sum_ {i, \ell} \mathbf {I} _ {i, j} ^ {\ell} \tag {12}
$$
is the sparsity loss, encouraging the model to prune the context. In total  $(Ln(n - 1)) / 2$  entries of  $\mathbf{I}_{i,j}^{\ell}$  are learned, as indicated in Eq. (6). We choose  $\gamma >0$  to enforce different levels of sparsity. In general, for a current position  $i$  in the context, we define as sparsity, the percentage of the previous tokens dropped, i.e. (tokens  $\leq i$  dropped)/i.