SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zichang Liu
Fangshuo Liao
Fangshuo.Liao@rice.edu
Victor Xie
vyx2@rice.edu
Anastasios Kyrillidis
anastasios@rice.edu
Aditya Desai
Aditya.P.Desai@rice.edu
Weitao Wang
wtwang@rice.edu
Zhaozhuo Xu
Large language models(LLMs) have sparked a new wave of exciting AI applications. Hosting these models at scale requires significant memory resources. One crucial memory bottleneck for the deployment stems from the context window. It is commonly recognized that model weights are memory hungry; however, the size of key-value embedding stored during the generation process (KV cache) can easily surpass the model size. The enormous size of the KV cache puts constraints on the inference batch size, which is crucial for high throughput inference workload. Inspired by an interesting observation of the attention scores, we hypothesize the persistence of importance: only pivotal tokens, which had a substantial influence at one step, will significantly influence future generations. Based on our empirical verification and theoretical analysis around this hypothesis, we propose SCISSORHANDS, a system that maintains the memory usage of KV cache under a fixed budget without finetuning the model. We validate that SCISSORHANDS reduces the inference memory usage of the KV cache by up to  $5 \times$  without compromising model quality. We further demonstrate that SCISSORHANDS can be combined with 4-bit quantization for further compression
Large language models(LLMs), trained on immense amounts of text data, have demonstrated an incredible ability to generate text that is both logically connected and contextually relevant [1-5]. LLM inference follows an autoregressive fashion, generating one token at each step conditioned on the previous steps. At each step, the key-value embedding in attention is stored in memory to avoid repetitive key-value projection computation at future steps. Unfortunately, the memory of the key-value cache (KV cache), including prompts and previously generated tokens, can be surprisingly large. Using OPT-175B as an example, the impressive 175 billion parameters consume around 325 GB of memory. At the same time, at batch size 128 and sequence length 2048, the KV cache requires around 950 GB of memory, three times larger than the model weights. Considering that 8 Nvidia A100-80GB offers 640GB GPU memory, the memory usage of the KV cache is truly concerning.
LLMs are typically deployed on fixed memory hardware, and the size of model weights is also fixed once deployed. Apart from a small memory buffer typically reserved for communication and computation, the rest of the available memory is for the KV cache. The size of the KV cache depends on batch size, sequence length, and model dimension. Thus, at a given inference sequence length, compression in the KV cache memory translates almost linearly into an increase in the batch size. And any increase in batch size is significant for high-throughput inference systems [6, 7].
Quantization and sparsity approaches [8-14] have been studied in LLMs to reduce the model sizes. However, compressing the KV cache remains an open but challenging problem. First, training models at the scale of hundreds of billions of parameters on a large amount of data is prohibitively expensive. Thus, an ideal compression algorithm should be applicable without training. Second, emerging applications such as dialogue systems require an extremely long context window. The maximum sequence length of LLMs is growing to over 32K [15]. The size of the KV cache also grows linearly with sequence length. For scalability, an ideal compression algorithm should reduce the memory from the sequence length dimension. At last, compression should preserve LLMs' quality and in-context learning ability.
(a) Attention map at position 178
(b) Attention map at position 228
(c) Attention map at position 278
Figure 1: Repetitive Attention Pattern. We plot the attention map at three token positions in a sentence. Only five attention heads are plotted for a clearer presentation. We discretize the attention score such that the high score is dark green, and the low score is light green. In Figure 1(a), the token at position 178 pays heavy attention to positions 27, 63, 98, etc. This pattern is also present in the attention maps of position 228 and position 278.
We go beyond the traditional model compression techniques to achieve such demanding requirements. We envision that not all tokens must be stored in memory for LLM to understand the context. Just like humans can skim through an article and grasp the main idea, LLMs may also be able to skim and comprehend. It is commonly observed that the attention score from one token follows a strong power law distribution [16-20], meaning that one token will only heavily attend to a small number of tokens. More importantly, we observe Repetitive Attention Pattern from different tokens in the sequence in a trained LLM( Figure 1). Certain tokens are more important throughout the paragraph. Specifically, for two different tokens, there are similarities between which tokens they are heavily attending to and similarities between which tokens they are ignoring.
Inspired by the above observation, we articulate the Persistence of Importance Hypothesis: Only pivotal tokens, which had a substantial influence at one previous step, will have a significant influence at a future step. This hypothesis, if true, suggests that it is possible to foresee which token is likely to be important for future generations. Fortunately, we empirically verify that later tokens in the
sentence mostly only attend to tokens that were heavily attended from the early tokens in a sentence. And the overlapping ratio is surprisingly high, over  $90\%$  in most of the transformer layers (Figure 2).
Based on the above two findings, we present SCISSORHANDS that exploits the persistence of importance hypothesis to realize LLM inference with a compressed KV cache. In Section 4, we present an efficient algorithm such that the size of KV cache is always less than a predetermined budget. A theoretical guarantee justifies that such a compressed KV cache can approximate the attention output. In Section 5, we empirically evaluate SCISSORHANDS and show that SCISSORHANDS reduces the memory usage of KV cache  $2 - 5 \times$  without compromising model quality. Reduction in the KV cache can directly result in a larger batch size. Further, we adopt quantization and show its compatibility with SCISSORHANDS.
2 Problem Description and Related Work
This paper considers the LLM inference workflow, specifically focusing on the memory usage for storing the keys and values in attention. Let  $d$  be the hidden dimension of the model,  $b$  be the batch size, and  $p$  be the length of prompt sentences. We are given the trained model weights,  $W_{K}^{i}\in \mathbb{R}^{d\times d}$ ,  $W_{V}^{i}\in \mathbb{R}^{d\times d}$  for the key and value projection matrix at the  $i^{th}$  transformer layer.
The standard LLM inference consists of two stages: prompting and token generation. In the prompting stage, the model takes the prompt sentences as the input, and the key/value embedding in attention is stored as a cache to reduce repetitive computation. Denote  $x_{\mathrm{prompt}}^i = [x_1^i,\dots,x_p^i]$ ,  $x_{\mathrm{prompt}}^i\in \mathbb{R}^{b\times p\times d}$  as the input to attention at the  $i^{th}$  transformer layer. Denote the key cache and value cache at layer  $i$  as  $\mathcal{K}^i$ ,  $\mathcal{V}^i\in \mathbb{R}^{b\times p\times d}$ ,  $\mathcal{K}_0^i = x_{\mathrm{prompt}}^i W_K^i$ ,  $\mathcal{V}_0^i = x_{\mathrm{prompt}}^i W_V^i$ .
In the generation stage, the model starts with the stored KV cache in the prompting stage and generates one token at each step. At each step, the KV cache gets updated. Given the input to attention at step  $t$  in the  $i^{th}$  transformer layer  $x_{t}^{i} \in \mathbb{R}^{b \times 1 \times d}$ .  $\mathcal{K}_{t+1}^{i} = [\mathcal{K}_{t}^{i}, x_{t}^{i}W_{K}^{i}]$ ,  $\mathcal{V}_{t+1}^{i} = [\mathcal{V}_{t}^{i}, x_{t}^{i}W_{V}^{i}]$ .
2.1 LLM Inference Memory Breakdown
In this section, we provide the memory consumption breakdown of LLMs. The memory footprint consists of three parts: model weights, KV cache, and activation buffer. The size of model weights depends on model configuration, such as the number of transformer layers and hidden size. The size of the KV cache depends on model configurations, sequence length, and batch size. The size of the activation buffer depends on parallelism strategy, model configurations, and implementation. The size of the activation buffer is considerably smaller than the previous two. As shown in Table 1, the size of the KV cache,  $2.5 \times -5 \times$  larger than model weights, can quickly become the bottleneck in memory consumption. At the same time, much research has been spent on extending the length of the context window. GPT-4-32K can process up to 32,768 tokens [15]. Longer sequence length would make the KV cache memory problem even more severe.
Assuming LLM generates until its maximum sequence length, we summarize the maximum batch size before going out of GPU memory on a box of 8 A100 80GB GPU in Table 2. At the GPT-3 scale with a maximum sequence length of 2048, batch size cannot exceed 35 without offloading. Small batch size limits the model inference throughput.
2.2 Efficient Attention
Computing the attention matrix necessitates a time complexity of  $O(n^{2})$ , where  $n$  is the sequence length. As a result, a line of work has been proposed to mitigate the computation burden of the attention mechanism [16-20]. These approaches exploit low-rank or sparsification to approximate the attention output. Besides, [21] realized exact efficient attention with wall-clock speed by optimizing
Table 1: The memory consumption of model weights and KV cache for three different LLMs at batch size 128 and sequence length 2048 shows that the KV cache dominates the memory consumption.
Model	# of Layer	Hidden Size	Weights (GB)	KV cache (GB)
OPT-175B	96	12288	325	1152
LLaMA-65B	80	8192	130	640
BLOOM	70	14336	352	950
Table 2: Maximum batch size before hitting out of memory on a box of 8 A100 80GB GPU when models are deployed with its maximum sequence length.
Model	OPT-175B	LLaMA-65B	BLOOM
Maximum Batch Size	34	102	36
the number of memory reads and writes. However, these approaches were evaluated mostly for training, focused on computation complexity, and did not address the KV-Cache memory usage introduced by auto-regressive language models.
Recently, there is active research attempting to apply quantization or pruning in LLM [8-14]. However, they mostly focus on reducing the size of model weights. Flexgen [7] applies quantization and sparsification to the KV cache; however, the memory of the KV cache is not reduced regarding sequence lengths. It stores the quantized KV cache for all tokens in CPU memory and loads all attention keys from CPU memory to compute attention scores. At the same time, methods such as Multi-Query-Attention(MQA) [22] change the attention design such that keys and values are shared across all attention heads. MQA requires training the model from scratch, while our works focus entirely on the inference stage.
3 The Persistence of Importance Hypothesis
We first present one interesting observation upon which the persistence of importance hypothesis is derived in Section 3.1. In Section 3.2, we discuss the hypothesis in detail with empirical verification. Then, in Section 3.3, we provide theoretical intuition on the reason behind such model behaviors.
3.1 Repetitive Attention Pattern.
Observation. We are interested in the attention score from the position  $t$  over all the words that come before it in the sentence. In Figure 1, we provide three attention maps of a sentence randomly drawn from the Colossal Clean Crawled Corpus (C4) [23] using OPT-6B. Each attention map is a discretized attention score calculated at a random position. We consider a score larger than  $\frac{1}{t}$  as significant as  $\frac{1}{t}$  indicates an averaging mixing score. High attention scores are marked with dark green.
Result. High attention scores are observed at the same set of tokens from various positions in the sentence. In all three plots, we see dark green at sequence positions 27, 63, 98, 121, 152, and 177, suggesting that these tokens received high attention at all three positions. We observe similar model behavior at different transformer layers with different text inputs. More plots are in Appendix A.
**Implication.** Even though small differences exist, repetitive attention patterns are evident in the attention maps. There exist specific tokens that keep receiving high attention. Meanwhile, these attention maps show sparsity: only a few tokens have high attention scores.
3.2 The Persistence of Importance Hypothesis
The repetitive attention pattern suggests that specific tokens are influential throughout the sequence. A stricter claim is that these tokens are the only ones that could be significant for a future step. Thus, we articulate the persistence of importance hypothesis.
The Persistence of Importance Hypothesis. With a trained autoregressive language model, only pivotal tokens, which had a substantial influence at one previous step, will have a significant influence at a future step.
If true, this hypothesis indicates the possibility of foreseeing what information in the previous sequences could be vital for future steps. This hypothesis is trivial when pivotal tokens include all tokens in the entire sentences. However, a much more interesting case is when pivotal tokens are a subset of previous words. This would enable us to reduce the size of the KV cache by throwing away the embedding of non-important tokens.
Pivotal Token. One natural indication of a token's influence is the attention score. We consider a token pivotal for position  $t$  if this token receives an attention score larger than threshold  $\alpha$  from the token at position  $t$ . Let  $S_{t}$  denote the set of pivotal tokens for position  $t$ .  $S_{a\rightarrow b}$  denote the set of
(a) Persistence Ratio
(b) Size of  $S_{0\rightarrow t}$
Figure 2: Persistence ratio and the corresponding size of the pivotal token set. The persistence ratio is over  $95\%$  in most layers, with decreases at the later layers. Meanwhile, the number of pivotal tokens is considerably smaller than the sequence length. This suggests that the pivotal tokens of later half sentences are almost all included in the set of first halves.
pivotal tokens for every position from  $a$  to  $b$ .
$$
S _ {a \rightarrow b} = \cup_ {t = a} ^ {t = b} S _ {t}
$$
Verification. We measure persistence ratio as an empirical test the hypothesis. Persistence ratio measures how many tokens in the pivotal token sets of the later part of the sentence are also in the pivotal token sets of the initial part of the sentence. Let  $l$  denote the length of the sentence. We record  $S_{1\rightarrow t} \in \{x_1,\dots x_t\}$ , tokens in  $\{x_{1},\ldots ,x_{t}\}$  who received high attention from every position until  $t$ . Then, we record  $S_{t + 1\to l} \in \{x_1,\dots,x_t\}$ , tokens in  $\{x_{1},\dots,x_{t}\}$  who received high attention from position after  $t$ . The persistence ratio is the intersection divided by the size of  $S_{t + 1\to l}$ . Formally,
$$
\text {P e r s i s t e n c e} \quad R a t i o = \frac {\left| S _ {t + 1 \rightarrow l} \cap S _ {0 \rightarrow t} \right|}{\left|\left\{x \mid x \in S _ {t + 1 \rightarrow l} , x \in \left\{x _ {1} , \dots , x _ {t} \right\}\right\}\right|}
$$
At the same time, we measure  $\frac{|S_{0\rightarrow t}|}{t}$ .  $|S_{0\rightarrow t}| = t$  indicates that every token substantially impacted at least one position, which is the trivial case of persistence of importance hypothesis. Our test is performed with OPT models [24] with different datasets such as OpenBookQA [25] and WikiText [26]. In our verification, we set  $t = \frac{l}{2}$ , which measures the overlapping between the first and later half of the sentences. Same as in Section 3.1, we set  $\alpha = \frac{1}{t}$ , which suggests an average score.
Result. We present our main results in Figure 2. First, given the current criterion of pivotal token and  $t$  value, the size of  $S_{0\rightarrow t}$  is considerably smaller than half of the sentence length. This verifies that we are not considering the trivial case of our hypothesis. Second, the persistence ratio is generally over  $95\%$ , with dips in the later transformer layers. The pivotal token set of the later half sentences is mostly included in the set of the first half sentences. Combining these two pieces of empirical evidence, we see positive evidence for our hypothesis test.
Implication. The hypothesis provides insights for understanding the behavior of LLMs and opens up new opportunities for reducing the KV cache memory. The hypothesis suggests the possibility of predicting the potentially influential tokens for future steps. The non-influential tokens are unnecessary to store in the memory, as they are unlikely to have high attention scores. This reduces the number of tokens stored in the KV cache and the computation required at the attention.
3.3 Attention Weights Decides the Pivotal Tokens
In the previous section, we verified that the significant tokens would continue to be significant. In this section, we try to understand the reasons for such phenomena. We consider the token generation process of a simplified model: a single-layer transformer model with single-head attention.
$$
x _ {t + 1} = \mathcal {F} \left(a _ {t}\right), \text {w h e r e} a _ {t} = \operatorname {s o f t m a x} \left(^ {1} / t \cdot x _ {t} W _ {Q} W _ {K} ^ {\top} X _ {t - 1} ^ {\top}\right) X _ {t - 1} W _ {V} W _ {O} \tag {1}
$$
$x_{t} \in \mathbb{R}^{1 \times d}$  is a row vector.  $X_{t-1} \in \mathbb{R}^{(t-1) \times d}$  denotes the aggregation of  $x_{1}, \ldots, x_{t-1}$ , where the  $j$ th row is  $x_{j}$ .  $W_{Q}, W_{K}, W_{V} \in \mathbb{R}^{d \times d}$  and  $W_{O} \in \mathbb{R}^{d \times d}$  are the attention weights. Lastly,  $\mathcal{F}: \mathbb{R}^{1 \times d} \to \mathbb{R}^{1 \times d}$  denotes the MLP block following attention block, a two-layer MLP with skip connections, given by
$$
\mathcal {F} (x) = x + W _ {2} \operatorname {r e l u} \left(W _ {1} x\right) \tag {2}
$$
Algorithm 1 Inference with Budget KV cache
Input: Memory Budget  $B$ , Maximum Sequence Length  $T_{\mathrm{max}}$   
Key Cache  $\bar{\mathcal{K}} \in R^{n \times d}$ , Value Cache  $\bar{\mathcal{V}} \in R^{n \times d}$ , where  $n = 0$   
while  $t < T_{\mathrm{max}}$  do  
    Model update  $\bar{\mathcal{K}}, \bar{\mathcal{V}}$  such that  $n \gets n + 1$   
    if  $n > B$  then:  
        Compress KV cache using Algorithm 2 such that  $n \leq B$ .  
    end if  
     $t \gets t + 1$   
end while
Algorithm 2 Compress KV Cache
Input: Key Cache  $\bar{\mathcal{K}}\in \mathbf{R}^{n\times d}$  , Value Cache  $\bar{\nu}\in R^{n\times d}$  , History Window Size w, Recent Window Size  $r$  , Drop Amount m, Generation Step t, Importance Record  $I\leftarrow \vec{0}\in R^t$    
for  $i\in [t - w,t]$  do  $\triangleright$  Consider tokens within history window  $I\gets I + \alpha_{i} <   \frac{1}{t}$ $\triangleright$  Increment the counter for low score token end for   
 $I[: - r]\gets 0$ $\triangleright$  Keep cache within the recent window   
Keep set  $S_{t}\gets \operatorname {Argsort}(I)[: - m]$    
Keep everything in  $S_{t}$  in  $\bar{\kappa}\in R^{n\times d}$ $\bar{\nu}\in R^{n\times d}$  such that  $n\gets n - m$
We are interested in the attention scores  $\alpha_{t} = \mathrm{softmax}(1 / t\cdot x_{t}W_{Q}W_{K}^{\top}X_{t - 1}^{\top})$ . Notice that  $\alpha_{t,j}$  scales with  $x_{t}W_{Q}W_{K}^{\top}x_{j}^{\top}$ . The following theorem characterizes the behavior of  $x_{t}W_{Q}W_{K}^{\top}x_{j}^{\top}$
Theorem 3.1. Let  $A = W_{V}W_{O}W_{Q}W_{K}^{\top}$  and let  $\lambda_{K}, \lambda_{Q}, \lambda_{V}, \lambda_{O}$  denote the largest singular values of  $W_{K}, W_{Q}, W_{V}, W_{O}$ , respectively. Consider the transformer in (1) with normalized inputs  $\| x_{t}\|_{2} = 1$  for all  $t$ . Let  $c, \epsilon > 0$  be constants. Assume that  $a_{t}x_{t + 1}^{\top} \geq (1 - \delta)\| a_{t}\|_{2}$  with  $\delta \leq \left(\frac{c\epsilon}{\lambda_Q\lambda_K\lambda_V\lambda_O}\right)^2$ . Then for all  $x_{\ell}$  satisfying  $x_{\ell}Ax_{\ell}^{\top} \geq c$  and  $x_{\ell}Ax_{\ell} \geq \epsilon^{-1}\max_{j\in [t], j \neq \ell}x_{j}Ax_{\ell}^{\top}$ , it holds that
$$
\frac {x _ {\ell} A x _ {\ell} ^ {\top}}{\| a _ {t} \| _ {2}} \left(\alpha_ {t, \ell} - 3 \epsilon\right) \leq x _ {t + 1} W _ {Q} W _ {K} ^ {\top} x _ {j} ^ {\top} \leq \frac {x _ {\ell} A x _ {\ell} ^ {\top}}{\| a _ {t} \| _ {2}} \left(\alpha_ {t, \ell} + 3 \epsilon\right) \tag {3}
$$
The proof is provided in Appendix B. Theorem 3.1 shows that under an assumption on the MLP in (2), for all  $x_{\ell}$  such that  $x_{\ell}Ax_{\ell}^{\top}$  is large enough,  $x_{t + 1}W_{Q}W_{K}^{\top}x_{j}^{\top}$  satisfies Equation (3). The assumption on the MLP  $a_{t}x_{t + 1}^{\top}\geq (1 - \delta)\| a_{t}\|_{2}$  essentially requires a large cosine similarity between the input and output of  $\mathcal{F}$ . This behavior can be empirically verified in Appendix A. Essentially, skip connection dominates the output because  $\| x\| _2\gg \| W_2\mathsf{relu}(W_1x)\| _2$ , resulting in a cosine similarity close to one between input and output. Equation (3) shows that despite a factor of  $\frac{x_{\ell}Ax_{\ell}^{\top}}{\|a_t\|_2}$ ,  $x_{t + 1}W_{Q}W_{K}^{\top}x_{j}^{\top}$  almost scales with  $\alpha_{t,\ell}$ . Since  $x_{t + 1}W_{Q}W_{K}^{\top}x_{j}^{\top}$  directly affects  $\alpha_{t + 1,\ell}$ , this property shows that a larger  $\alpha_{t,\ell}$  will potentially imply a large  $\alpha_{t + 1,\ell}$ .
Our theorem shows that the property in Equation (3) property only holds for  $x_{\ell}$  such that  $x_{\ell}Ax_{\ell}^{\top}$  is large.  $A$  are trained attention weights. This condition may suggest that the trained weights  $A$  selects  $x_{\ell}$  as a pivotal token. Each attention is learned to identify some subspace. Only those tokens embedded inside these regions are pivotal for this attention. This would explain why only some specific tokens are always relevant.
4 Sequential Token Generation Under budget
In this section, we present SCISSORHANDS, which reduces the KV cache memory from the sequence length dimension without fine-tuning the model. In Section 4.1, we describe how SCISSORHANDS maintains the KV cache under a given budget. Section 4.2 provides a theoretical analysis of the algorithm and the approximation error.
4.1 Budget KV Cache for Single Attention Head
In this section, for the sake of the discussion, we drop the layer number notation  $i$  and batch size dimension.  $\mathcal{K}_t,\mathcal{V}_t\in R^{t\times d}$  denote for the KV cache until step  $t$ .  $x_{t}\in \mathbb{R}^{1\times d}$  is a row vector that denotes the input to attention at step  $t$ . The output of an attention head at step  $t$  can be written as,
$$
a _ {t} = \sum_ {i = 1} ^ {t} \alpha_ {t, i} \mathcal {V} [ i ] _ {t}, \text {w h e r e} \alpha_ {t, i} = \frac {\exp \left(\langle x _ {t} W _ {Q} , \mathcal {K} _ {t} [ i ] \rangle\right)}{\sum_ {i = 1} ^ {t} \exp \left(\langle x _ {t} W _ {Q} , \mathcal {K} _ {t} [ i ] \rangle\right)}
$$
Intuition. As shown in Section 3, the attention scores  $\alpha_{t,i}$  follow a strong power-law distribution. For the autoregressive generation process, if there exists an oracle such that we can identify the heavy score tokens before the future generation step, then the memory of the KV cache can be significantly reduced by only storing the heavy score tokens. Fortunately, the persistence of importance hypothesis provides us with such an oracle. It states that only historical tokens with significant contributions toward previous generated tokens will have significant contributions toward future tokens.
Challenges. LLMs are deployed on hardware with a fixed memory. The algorithm should maintain the cache under fixed memory to meet the hard requirement. Further, LLMs are already computationally intensive. The algorithm should avoid introducing much extra burden on computation.