A Simple and Effective L\_2 Norm-Based Strategy for KV Cache Compression (2024)

Abstract: The deployment of large language models (LLMs) is often hindered by the extensive memory requirements of the Key-Value (KV) cache, especially as context lengths increase. Existing approaches to reduce the KV cache size involve either fine-tuning the model to learn a compression strategy or leveraging attention scores to reduce the sequence length. We analyse the attention distributions in decoder-only Transformers-based models and observe that attention allocation patterns stay consistent across most layers. Surprisingly, we find a clear correlation between the L_2 norm and the attention scores over cached KV pairs, where a low L_2 norm of a key embedding usually leads to a high attention score during decoding. This finding indicates that the influence of a KV pair is potentially determined by the key embedding itself before being queried. Based on this observation, we compress the KV cache based on the L_2 norm of key embeddings. Our experimental results show that this simple strategy can reduce the KV cache size by 50% on language modelling and needle-in-a-haystack tasks and 90% on passkey retrieval tasks without losing accuracy. Moreover, without relying on the attention scores, this approach remains compatible with FlashAttention, enabling broader applicability.

## introduction handling long contexts is desirable for large language models (llms), as it allows them to perform tasks that require understanding long-term dependencies @xcite @xcite @xcite . a key component for modelling long context is the kv cache, which stores the keys and values of past tokens in memory to avoid recomputing them during generation. however, processing long-context inputs often results in a high decoding latency since it requires repeatedly reading a potentially large kv cache from high-bandwidth memory (hbm) to the streaming multiprocessor (sm) during decoding @xcite . consequently, the practical deployment of llms is frequently hindered by hardware limitations. to address the issue of kv cache growth, various kv cache compression methods have been proposed. these methods can be broadly categorised into trainable approaches, which involve modifications to the model architecture @xcite , or fine-tuning regime to inherently manage kv cache size @xcite , and non-trainable approaches, which apply post-hoc compression techniques to reduce the cache footprint without altering the underlying model @xcite . while these methods have shown promise, they often involve complex algorithms or significant computational overhead, limiting their practicality; for example, post-hoc compression algorithms usually evict kv pairs based on attention scores, which is not compatible with flashattention @xcite and thus prevents their we show that the l 2 norm of cached keys has a high correlation with attention scores. more specifically, we observe that a low l 2 norm of a key embedding usually leads to a high attention score during decoding. based on this observation, we propose a simple and highly effective strategy for kv cache compression: keeping in memory only the keys with lowest l 2 norm, and the corresponding values. unlike many existing methods, our heuristic can be applied off-the-shelf to any transformer-based decoder-only llm without the need for additional training or significant modifications. more importantly, our method estimates the influence of cached key-value pairs without the need to compute the attention scores. therefore, unlike other compression methods @xcite , it can be easily integrated with the popular flashattention @xcite . our experimental results demonstrate that this heuristic allows maintaining model performance in language modelling tasks and in tasks that require the model to store and retrieve the most critical information, such as passkey retrieval @xcite and needlein-a-haystack tasks @xcite . ## patterns in the attention matrix we first examine the attention scores on the language modelling task for a range of popular llms. by analysing the key embeddings and the attention distribution, we observe that key embeddings with low l 2 norm are often associated with higher attention scores. in fig. 2 , we provide an example using llama-2-7b @xcite , where the first row presents the attention distribution over the kv pairs, and the second row presents the l 2 norm of each key embedding. we observe that the tokens with high attention scores, such as "<@xmath0.", have significantly lower l 2 norm values than others. while @xcite already observed peaked attention distributions for specific tokens, and @xcite pointed out the influence of high l 2 norm tokens on attention maps, we are the first, to the best of our knowledge, to point out the correlation between the l 2 norm of the key embeddings and attention score. based on our observation, we consider the following research question: can we compress the kv cache based on the l 2 norm of the key embeddings? an intuitive way to estimate the influence of compressing the kv cache is by examining the attention scores that are dropped due to the compression. in the following, we formally define this influence. given a prompt consisting of n tokens (x 1 , x 2 , ..., x n ), the llm first encodes them into a kv cache-this step is referred to as the pre-filling phase. then, the model autoregressively generates the next token x n+1 . when performing kv cache compression, some key-value pairs may be dropped and thus cannot be attended to. we define the attention loss caused by the compression as the sum 0 1000 2000 3000 4000 5000 6000 input length 1 2 3 4 5 6 7 log ppl llama 2-7b (wikipedia) no_compression max kv 2000 (keep low norm) max kv 2000 (keep random) max kv 2000 (keep high norm) 0 2000 4000 6000 8000 input length 1 2 3 4 log ppl llama 3-8b (wikipedia) no compression max kv 2000 (keep low norm) max kv 2000 (random) max kv 2000 (keep high norm) max kv 2000 (fastgen) 0 2000 4000 6000 8000 input length 2 4 6 8 log ppl gemma (wikipedia) no compression max kv 2000 (keep low norm) max kv 2000 (random) max kv 2000 (keep high norm) of the attention scores associated with the dropped kv pairs: we name the y l,h as alr, which indicates the attention loss for a compression method using ideal attention loss as reference. in fig. 1 , we plot the y across layers and heads. we observe that heads in the first two layers and some middle layers around the 12th layer have relatively high y values. the heads in other layers have lower y values, indicating a high correlation between l 2 norm and attention score. by leveraging this correlation, we can compress the kv cache based on the l 2 norm of key embeddings. optionally, we can skip the compression at the layers with low correlation. we show ablation experiments skipping layers in appendix b. ## experiments we evaluate our method on language modelling and two long-context modelling tasks, i.e., needlein-a-haystack and passkey retrieval. based on the observation supported by fig. 1 , the heads in the first two layers usually have a low correlation between l 2 norm and attention score, so we do not perform compression on these layers as default. we conduct experiments to investigate the impact of compression on different layers in appendix a. language modelling for language modelling, we let the kv cache grow until a specific predefined length and subsequently start to discard the tokens with the highest l 2 norm. we show in fig. 3 that evicting even up to the 50% of kv cache does not impact perplexity. perplexity increases, as expected, once we exceed the pre-training context length. we show more results, including next token accuracy in appendix b. to further verify that keys with low l 2 norm capture significant information, we test other eviction strategies, i.e. keeping tokens with highest l 2 norm and keeping random tokens. it is clear from fig. 3 that discarding tokens with low l 2 impairs performance, even more so than random discarding, thus highlighting the importance of these low l 2 norm keys. pressure test on long-context tasks the needle-in-a-haystack task (kamradt, 2023) and passkey retrieval task (mohtashami and jaggi, 2023) are two synthetic tasks that are widely used 10% 30% 50% 70% to pressure test the long-context modelling capability of llms. in both tasks, the model needs to identify and retrieve the important information from a long context to generate correct answers. thus, these tasks test the compression method's ability to keep important kv pairs and drop redundant ones. in figure 4a and figure 4b , we present the experimental results of llama-2-7b-80k @xcite . we analyse additional models in appendix c. the model retains its performance on the needle-in-ahaystack task with 30% kv cache compression (figure 4a ) and maintains 99% accuracy with 50% compression. it also achieves 100% accuracy on the passkey retrieval task, even with 90% kv cache compression (figure 4b ). we compare other eviction strategies, like keeping kv pairs with low l 2 norm, with high l 2 norm, and random. in figure 4a and figure 4b , we observe that the model cannot answer correctly when keeping only high l 2 norm kv pairs, obtaining near zero accuracy. when we randomly compress the kv cache, the performance decreases significantly faster than keeping low l 2 norm kv pairs. the above analysis indicates that kv pairs with low l 2 norm are critical to generating the correct answer and thus contain important information. experiments on longbench additionally, we evaluate on longbench @xcite . we test on several subsets, including narrativeqa (kociský et al., 2018), qasper (dasigi et al., 2021), hotpotqa @xcite , 2wikimqa (ho et al., 2020), and qmsum @xcite . we report the results for the recently released long context llama3.1 in figure 4c . in addition, we show the complete per-subset results in appendix c. the experimental results show that compressing the kv cache with low l 2 norm only introduces a small ac-curacy decrease even when compressing 50% kv cache, while compressing kv cache with high l 2 norm results in almost zero accuracy. comparison with fastgen like the majority of methods in the literature, fastgen (holmes et al., 2024) utilises attention scores, which makes it incompatible with the popular flashattention @xcite , thereby limiting its efficiency and usability. for a fair comparison, we implement fastgen without using the attention scores, i.e., we only consider local, punctuation and special tokens. we perform experiments on language modelling with the llama3 model @xcite . our method still outperforms fastgen with up to 50% kv cache eviction. we show the results in figure 5 . relationship between embedding and l 2 norm after identifying a correlation between the l 2 norm of token key embeddings and attention scores, we performed a further exploration by analyzing the key prjections in the kv cache. we found that tokens with lower l 2 norm show sparse activations, with few dimensions having high values while most remain near zero, indicating limited use of the vec0 5 10 14 12 10 8 6 4 2 0 0 5 10 14 12 10 8 6 4 2 0 0 5 10 14 12 10 8 6 4 2 0 0 5 10 14 12 10 8 6 4 2 0 head 0 head 11 head 23 head 31 (a) attention maps of llama3-8b when zeroing out random activations. tor space (figure 6 ). this sparsity aligns with the concept of "sink" tokens @xcite , where many queries align with certain tokens, increasing their attention scores. we hypothesise that the lower l 2 norm reflects a partial use of the available embedding space, leading to increased attention for these tokens. to test this, we zeroed out the dimensions responsible for the peaked activations in low l 2 norm key embeddings and observed significant changes in attention maps during generation (figure 7 ). however, randomly altering dimensions did not produce the same effect. this finding suggests that the l 2 norm may serve as a proxy for the extent to which an embedding utilises the available vector space and, consequently, the degree to which it influences attention. lower l 2 norm appears to correspond to embeddings that drive disproportionately high attention values due to their alignment with a common "sink" direction. recent long-context llms like gemini-pro-1.5 (reid et al., 2024), claude-3 (anthropic, 2024), and gpt4 (achiam et al., 2023) can process hundreds of thousands of tokens, but face high decoding latency. to address this, works like pageattention (kwon et al., 2023), infinite-llm (lin et al., 2024), and vattention (prabhu et al., 2024) propose efficient memory management strategies. others focus on kv cache compression: dmc (nawrot et al., 2024) uses dynamic trainable token merging, while h2o (zhang et al., 2024b), fastgen (ge et al., 2023), and snapkv (li et al., 2024) employ various attention-based training-free compression strategies. unlike these methods, we uniquely use key embedding l 2 norm for compression. while (darcet et al., 2024) had previously found that high l 2 norm hidden states aggregate important information, we are the first, to the best of our knowledge, to discover and leverage the correlation between low l 2 norm key embeddings and high attention scores for efficient kv cache compression. ## conclusion we introduced a simple yet highly effective strategy for kv cache compression in llms based on the l 2 norm of key embeddings. we show that there is a significant correlation between the l 2 norm of a key embedding and its attention score. leveraging this observation, we compress the kv cache by retaining only those keys with the lowest l 2 norm. our experimental results on various tasks show that our compression strategy maintains the predictive accuracy of the model while significantly reducing the memory footprint. our approach is straightforward and can be applied directly to any transformer-based, decoder-only llm.