Sequence can Secretly Tell You What to Discard
Jincheng Dai Zhuowei Huang Haiyun Jiang Chen Chen Deng Cai Wei Bi Shuming Shi



Input-Output Specifications of Grid-Forming Functions and Data-Driven Verification Methods
    Jennifer T. Bui and Dominic Groß This work was supported by the U.S. Department of Energy's Office of Energy Efficiency and Renewable Energy (EERE) under the Solar Energy Technologies Office Award Number 38637. The views expressed herein do not necessarily represent the views of the U.S. Department of Energy or the United States Government. J. T. Bui and D. Groß is with the Department of Electrical and Computer Engineering at the University of Wisconsin-Madison, USA; e-mail:{jtbui, dominic.gross}@wisc.edu
    April 29, 2024
==================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================





  Large Language Models (LLMs), despite their impressive performance on a wide range of tasks, require significant GPU memory and consume substantial computational resources. 
  In addition to model weights, the memory occupied by KV cache increases linearly with sequence length, becoming a main bottleneck for inference. In this paper, we introduce a novel approach for optimizing the KV cache which significantly reduces its memory footprint. Through a comprehensive investigation, we find that on LLaMA2 series models, (i) the similarity between adjacent tokens' query vectors is remarkably high, and (ii) current query's attention calculation can rely solely on the attention information of a small portion of the preceding queries. Based on these observations, we propose CORM, a KV cache eviction policy that dynamically retains important key-value pairs for inference without finetuning the model. We validate that CORM reduces the inference memory usage of KV cache by up to 70% without noticeable performance degradation across six tasks in LongBench.
  





§ INTRODUCTION


Large language models (LLMs) have demonstrated impressive proficiency in a wide range of natural language processing tasks such as question answering, summarization and multi-turn dialogues <cit.>. 
Considering substantial cost of deploying LLMs introduced by tremendous model size and quadratic cost of attention layer, many works focused on model compression and memory-efficient attention techniques <cit.>. 
However, the size of KV cache, which stores previous tokens' key and value states to avoid re-computation, scaling linearly with sequence length during generation, also incurs significant overhead. For instance, even a 7 billion-parameter model with batch size of 128 and sequence length of 4096 results in 256GB of KV cache, far exceeds memory consumed by model itself which is only 14GB. A natural idea is to discard some less informative KV cache to reduce memory consumption. The challenge lies in finding a balance between discarding as much as possible while still maintaining model performance.

Despite multi-query attention <cit.> and grouped-query attention <cit.> can reduce the size of KV cache by reducing attention heads, it needs re-training to recover performance of original model. Recent works <cit.> have investigated implementing KV cache using specific eviction policy, that determines which key-value states should be evicted from KV cache. These methods aim to compress KV cache to a pre-defined budget size, thereby reducing memory and computational overhead. However, they save same number of key-value pairs for all attention heads and layers, ignoring that the number of keys playing an important role may vary across different attention heads and layers <cit.>. 



Intuitively, if important information in the KV cache exceeds the predetermined budget size, the performance of the model is likely to decline as it unavoidably evicts some crucial information. Our preliminary exploration also reveals that different attention layers and heads show different sparsities as shown in Figure <ref>. First, we observe that bottom layers of the model are relatively dense[Let t denote sequence length, we count the proportion of keys which attention score larger than average score 1/t and denote it as r. The larger r is, the sparser the layer is. ], while the remaining attention layers exhibit significant sparsity. Second, even within the same layer, different heads can exhibit obvious differences in sparsity levels. These properties suggest that we need to treat different layers and heads differently, rather than using the same budget size for all of them. In addition, we prove that completely similar queries have similar concerns about keys, and observe that recent query vectors are quite similar on LLaMA2 series models so current query can directly use recent query attention messages during generation.

Based on the above insights, we first define the generation process of LLMs with a budget-unrestricted KV cache in Section <ref>. 
Then we propose Cache Optimization with Recent Message (CORM), a framework that exploits recent query attention information for KV cache optimization and token generation of LLMs. Specifically,



  * In Section <ref>, we explore the similarity between query vectors of all tokens within same sequence, revealing that recent query vectors are highly similar, which implies that (i) keys that are important for recent queries might be also important for the current query; and (ii) removing key-value pairs that appear to be less informative for recent queries can greatly preserve the performance of the model.


  * In Section <ref>, we present a simple method which dynamically evicts minor key-value pairs determined by recent tokens' attention information.


We conduct extensive experiments on LLaMA2-7B-Chat, considering its popularity and wide usage, to evaluate CORM across 6 tasks from LongBench <cit.> containing question answering, summarization, code completion, etc. 
Experiments show that even without explicitly setting a budget size, our method is still possible to achieve a high compression rate. Our method achieves better performance compared to StreamingLLM <cit.>, Scissorhands <cit.> and H_2O <cit.> with over 70% KV cache reduction rate and can even come close to fully restoring the performance of the model. 




§ RELATED WORK




  
Attention Let x ∈ℝ^n × d denote the input embeddings from a sequence of n feature vectors of dimension d. The multi-head self-attention <cit.>, as a core module of Transformer model, facilitates contextual information interaction within each head in the following manner:

    Q    = xW_q,
    
        K    = xW_k,
    
        V    = xW_v,
    Attention(x)    = softmax(QK^T/√(d_h)) × V


Q, K, V represent the query, key, and value matrices, which are obtained by linearly mapping x using weight matrices W_q, W_k, and W_v ∈ℝ^d × d_h, respectively. d_h is the dimension of each individual head. 



  
KV Cache According to autoregressive paradigm, transformer decoder model predicts future tokens based on both previous and current tokens. Recalculating the key-value pairs for previous tokens at each decoding step is clearly an inefficient strategy. A common practice is to retain the key-value pairs of previous tokens for subsequent reuse. Thus, the consumption of KV cache becomes linearly correlated with the length of input sequence. When dealing with long contexts, however, the use of such a space-time trade-off approach may still pose challenges.



  
Training Policies The advent of multi-query attention (MQA) <cit.> is to address the influence of attention heads on KV cache within multi-head attention (MHA) mechanism. It facilitates the sharing of the same set of keys and values among different heads to alleviate cache pressure. Grouped-query attention (GQA) <cit.> represents a trade-off between MHA and MQA, achieving key-value sharing within each group through mean-pooling-based uptraining. Both methods require additional training to restore model performance due to the inability to directly convert.



  
Training-free Policies During generation, sequence length is the primary factor of cache pressure. Recent methods aim to balance model efficiency and inference cost without extra training and architectural changes. StreamingLLM <cit.> keeps attention sink token and recent tokens throughout decoding process to align with the training window. Scissorhands <cit.> maintains pivotal tokens and recent tokens based on the persistence of importance hypothesis. H_2O <cit.> utilizes accumulated attention score to maintain heavy hitters and recent tokens. TOVA <cit.> removes tokens with the lowest current attention score from the fixed cache at each decoding step. RoCo <cit.> retains tokens in the fixed cache based on high mean cumulative attention scores and top r standard deviations. Aforementioned methods consistently operate on a fixed cache, ignoring that the number of tokens playing an important role may vary across different attention heads and layers. 



§ OBSERVATIONS


We first demonstrate the existence of attention sparsity in LLMs in Section <ref>, then discuss the phenomenon that similar queries have similar attention concerns for keys in Section <ref>. In Section <ref>, we show an intriguing observation that current query is most similar to recent queries.



 §.§ Attention sparsity in LLMs



We first explore the sparsity in attention layers of LLMs, which provides an effective guarantee for us to reduce KV cache size. Specifically, we use proportion of important keys to represent attention sparsity. Let q_t∈ℝ^1 × d denote the query state vector at step t, k_i∈ℝ^1 × d denote the key state vector at step i (1 ≤ i ≤ t), where d is hidden dimension (for the sake of simplicity, we only consider a single head here). The normalized attention score of q_t for k_i is computed as:


    α_t,i = exp(q_t k^T_i / √(d))/∑_j=1^texp(q_t k^T_j / √(d)).




  
Definition 3.1 (Important Key) We define a key k_i is considered important in step t, if and only if α_t,i≥1/t, otherwise it is considered minor.

We conduct zero-shot inference with LLaMA2-7B model on the test set of PG-19 <cit.>. We plot the layer-wise and head-wise sparsity within attention blocks, the results are presented in Figure <ref>. It reveals that bottom layers are relatively dense, while other layers are highly sparse with over 90% sparsity.
This makes it possible to do attention computation on only small part of KV cache during generation.




 §.§ Similar queries have similar concerns for keys



The previous section reveals the existence of attention sparsity in LLMs, which provides an opportunity to reduce KV cache size while maintaining performance. In this section we give a theoretical analysis that similar queries have similar concerns for keys for eviction policy design.

Consider the i-th and j-th query state vectors q_i and q_j in a sequence of token length T (i < j ≤ T). Their cosine similarity can be computed as:

    cosine_similarity(q_i, q_j) = q_i q_j^T/ q_i· q_j.


Consider all key states k_1, k_2, ..., k_i-1 before i-th key. Assume that cosine_similarity(q_i, q_j)=1, then q_i=m · q_j with m ∈ℝ^+. The attention weight[attention weight is unnormalized attention score] of q_i to the previous i-1 keys can be represented as:


    attention_weight = 1/√(d) (q_ik_1^T, q_ik_2^T, ..., q_ik_i-1^T) = m/√(d)· (q_jk_1^T, q_jk_2^T, ..., q_jk_i-1^T).


Note that m is a positive number that does not affect the relative order of the attention weights. For example, for q_i, if q_ik_1^T > q_ik_2^T, there must be q_jk_1^T > q_jk_2^T for q_j. 
This means if a key is important to q_i, it is also important to q_j, though the degree of importance may vary due to the softmax function.



Although it's nearly impossible that cosine_similarity(q_i, q_j)=1 in real situation, we can make the hypothesis that two similar queries may have similar concerns for keys. 
To validate this hypothesis, we provide two attention maps of a sentence randomly drawn from PG-19 using LLaMA2-7B, as shown in Figure <ref>. Important keys are marked with bright green, more plots are available in Appendix <ref>. 
We observe that the hypothesis is true, and similar queries exhibit similar concerns for important keys. At the same time, important keys only account for a small proportion especially in deeper attention layers, which is consistent with the finding that deeper layers are sparser in previous section.



 §.§ Similarity exploration of query vectors


We have validated two similar queries have similar concerns for keys in Section <ref>, we also need to validate that at each step we can find a previous query state that is similar enough to current query state in same layer and same head. To check this, we visualize cosine similarity of query vectors within same sequence as shown in Figure <ref>, more plots are available in Appendix <ref>. We observe an intriguing phenomenon that many images show clear oblique color segmentation, with the top oblique block closest to dark red which means current query is most similar to recent queries.



Through above observations, we see an opportunity to design a KV cache eviction policy based on query similarity that preserves the LLM generation performance. 



§ CACHE OPTIMIZATION WITH RECENT MESSAGE



In this section, we present CORM, a method reduces the KV cache memory based on recent query attention information without any fine-tuning process. In Section <ref>, we derive that current query can directly use recent query attention messages during generation. In Section <ref>, we present CORM eviction policy and describe how it works during generation. 




 §.§ Generate based on recent query attention messages


Consider observations in Section <ref>, intuitively, we can directly store all queries and their attention information for future reference. At each generation step, use current query to find the most similar one from previous queries, and use its saved attention information to calculate solely on important keys. However, this approach incurs a significant cost. First, storing all queries results in a substantial increase in memory overhead. Second, the requirement of performing similarity calculations at each step adds to the computational overhead.

Since in most cases current query is most similar to recent queries as described in Section <ref>, we can just use recent query attention messages. And from Figure <ref> we can also observe that only a small proportion of keys are considered important by recent queries. Therefore even if we save all the keys that are considered important in previous steps, we can save a lot of memory. 



 §.§ Eviction algorithm via recent message


We have shown recent query attention information is enough for cache optimization in Section <ref>. In the following, we formally define this algorithm and introduce how to integrate it into LLM generation directly.



  
Definition 4.1 (Long-term Minor Key) A key k_i is considered as long-term minor key only if it is considered minor in all recent r steps (from t-r+1 to t).



  
Approach CORM will have a recent window of size w to record the information of recent w queries, and will always keep recent r keys unremoved to prevent them from being discarded prematurely due to insufficient observations. During generation, k_i, v_i will be discarded once k_i is regarded as long-term minor key. For better explanation we present pytorch code[For the sake of brevity, the code snippet only demonstrates single-head eviction operation, while in the actual implementation, it will be performed on each head at every layer.] of main algorithm in Algorithm <ref>. Intuitively, when w is larger, more keys and values will be saved, the compression rate will be smaller and performance will be better; Conversely, when w is smaller, fewer keys and values will be saved, the compression rate will be larger and performance will be worse. So there's a tradeoff between performance and compress rate.





  
Memory Overhead Analysis In order to reduce memory overhead of KV cache, an extra memory overhead is introduced by recent information cache. We need to store recent query messages which increase memory overhead. However, these overheads are far less than compressed KV cache, one can use a small portion of memory to avoid maintaining full KV cache memory without obvious performance degradation. On the other hand, the compression rate will increase as the sequence length increases as shown in Figure <ref>, resulting in a lower memory overhead for this component in comparison.











§ EMPIRICAL EVALUATION


In this section, we present the results that demonstrate CORM can reduce up to 70% of the memory footprint of KV Cache without accuracy degradation on LLaMA2-7B-Chat. 




  
Dataset To broadly validate feasibility of our method on real-world use cases, we choose LongBench <cit.> as our evaluation benchmark, which contains a wide range of long-text tasks such as question answering <cit.>, summarization <cit.>, few-shot learning <cit.>, synthetic task and code completion <cit.>. Here we do not consider short text tasks, because even full cache doesn't have any bottlenecks.



  
Models Since sequence length is the main factor in the continuous growth of KV Cache, we employ LLaMA2-7B-Chat <cit.> for 4K test considering its wide usage.



  
Baselines Since CORM reduces KV cache without need for training, we consider several similar approaches as our baselines: StreamLLM <cit.>, Scissorhands <cit.> and H_2O <cit.>. In addition, the full KV cache is also considered as strong baseline to measure the performance loss of other methods.



  
Setting All baselines can be regarded as fixed budget size KV cache compression, however CORM is a dynamic compression method. Since we find that CORM has similar compression rates for various task texts with the same sequence length. For fair comparison, we plot the relationship between model compression rate and sequence length using texts randomly sampled from PG19 <cit.> as shown in Figure <ref>.  





  
Main Results We evaluate LLaMA2-7B-Chat for 4K length text. Results are summarized in Table <ref> & <ref> for LLaMA2-7B-Chat. The following observations can be drawn: (1) CORM consistently outperforms previous methods at the same compression rate across a wide range of tasks. (2) Meanwhile, with over 70% KV cache reduction, CORM achieves comparable performance as the model with full KV cache and even surpass it on some tasks, we speculate it's because there's some noise in full KV cache that affects model output and our method can eliminate this noise to a certain extent by discarding some KV cache. 







 §.§ Budget unnecessity: is unbudgeted better?

We primarily focus on the effectiveness of not setting a budget versus setting a fixed budget. Note that since we use same window size and recent size as Scissorhands in the experiment, it can be regarded a natural ablation experiment. And Table <ref> & <ref> have shown that, at the similar compression rate, CORM is much better than Scissorhands in most tasks, and performance of other tasks is close. This verifies that different transformer layers and heads should be treated differently rather than setting a same fixed budget size.





§ CONCLUSION

In this paper, we investigate a critical memory bottleneck in LLM deployment, KV cache. Inspired by similar queries have similar concerns for keys and recent queries are similar enough, we propose CORM, an unbudgeted KV cache eviction policy for significantly reducing its memory footprint, by reusing recent query attention messages. Through extensive evaluations, we demonstrate that CORM can reduce the inference memory usage of the KV cache by up to 70% without noticeable performance degradation across a variety of tasks.




unsrtnat








§ MORE PLOTS





 §.§ Similar queries share most of important tokens


We provide the attention map similar to Figure <ref> but from different heads on the same text in Figure <ref>, Figure <ref>. Plots from a different layer on the same text are shown in Figure <ref>.









 §.§ Relationship of query vectors


We provide the query vectors' cosine similarity visualization similar to Figure <ref> but from different layers and heads on the same text in Figure <ref>.





§ TASK MAPPING


An overview of the dataset statistics and mapping from ID to dataset in LongBench is shown in Figure <ref>.



