RazorAttention: Efficient KV Cache Compression Through Retrieval Heads

## Abstract

Abstract The memory and computational demands of Key-Value (KV) cache present significant challenges for deploying long-context language models. Previous approaches attempt to mitigate this issue by selectively dropping tokens, which irreversibly erases critical information that might be needed for future queries. In this paper, we propose a novel compression technique for KV cache that preserves all token information. Our investigation reveals that: i) Most attention heads primarily focus on the local context; ii) Only a few heads, denoted as retrieval heads, can essentially pay attention to all input tokens. These key observations motivate us to use separate caching strategy for attention heads. Therefore, we propose RazorAttention, a training-free KV cache compression algorithm, which maintains a full cache for these crucial retrieval heads and discards the remote tokens in non-retrieval heads. Furthermore, we introduce a novel mechanism involving a “compensation token” to further recover the information in the dropped tokens. Extensive evaluations across a diverse set of large language models (LLMs) demonstrate that RazorAttention achieves a reduction in KV cache size by over 70% without noticeable impacts on performance. Additionally, RazorAttention is compatible with FlashAttention, rendering it an efficient and plug-and-play solution that enhances LLM inference efficiency without overhead or retraining of the original model.

## 1 Introduction

Long-context large language models (LLMs) have significantly advanced capabilities in natural language processing across diverse tasks. However, the growth of the Key-Value (KV) cache under increasing input length has become the major bottleneck for deployment. There are been plenty of previous work designed to alleviate this problem by compressing the KV cache size, including quantization  [, , ], token-dropping  [, ], local attention  [, ], etc.

Figure: Figure 1: RazorAttention achieves comparable performance to the original model, even with $70\%$ KV cache compressed. To demonstrate this, we tested Llama2-13B-64K [] on the Needle in A Haystack benchmark [].
Refer to caption: x1.png

Figure: Figure 2: Importance-based token-dropping methods cannot work when querying the less relevant information to the main theme. Here, we use an 8K document from LongBench [] and add two sentences that are not relevant to the main theme. In this case, H2O discards tokens that are less relevant to the main theme, leading to failures in both Q1 and Q2. SnapKV discards tokens based on the first query, making it effective for Q1 but failing in subsequent queries like Q2. Only RazorAttention successfully outputs the exact information from the lengthy input even when we compress $70\%$ of the KV cache.
Refer to caption: x3.png

One major direction for KV cache compression is to directly drop tokens deemed unimportant so far [, , , ].
These methods inherently assume that tokens considered unimportant will not be needed in future queries, which does not hold in practical scenarios. For instance, a user might request information that is not directly aligned with the main theme of the processed text, or engage in a multi-round conversation querying different segments from the context. In these cases, the importance-based token-dropping methods can lead to significant performance degradation since the actual information required by the query might be discarded if considered unimportant (see our example on Qwen1.5-7B-Chat [] in Figure 2). This leads us to pose a critical question:

“Can we find a way to reduce the KV cache size without losing semantic information?”

In this work we address this problem from a novel perspective. Our investigation reveals that there exists a “retrieve and process” mechanism in LLMs when processing a long context. More specifically, LLMs can accurately recall the queried information from a lengthy input through certain group of attention heads, which we denote as “retrieval heads” (see Section 3.3 for definition). These heads are capable of concentrating most of their attention weights on the relevant information (w.r.t. the queries) and increasing the output probability for those words.
Another important finding is that non-retrieval heads primarily focus on local context or the attention sink [], which means these heads cannot effectively utilize all the semantic information from the input.
Based on these important findings, we hypothesize that LLM runs the reasoning procedure on a “retrieve and process” basis. That says, the model first uses the retrieval heads to gather relevant information, and then non-retrieval heads to process the retrieved information and generate the final response. This motivates us to design separate caching strategies for different heads: For retrieval heads, we keep the KV cache unaltered; for the rest heads, we only cache recent tokens and attention sinks.

Beyond this, we notice that there still exists a certain accuracy gap when we directly discard all the remote tokens in the non-retrieval heads. Therefore for these non-retrieval heads, we designed a “compensation token” for compressing the dropped cache into one token, and proved that the accuracy degradation due to the truncated KV cache gets further improved with this compensation token. With retrieval heads and compensation tokens, we prove that our algorithm, namely RazorAttention, can successfully compress $70\%$ of the KV cache without noticeable performance degradation as illustrated in Figure 1.

Last but not least, previous importance-based token-dropping methods cannot be combined with FlashAttention due to their reliance on the attention weights to compute the importance score, making them impracticable for implementation since FlashAttention is one of the most important components in long-context inference. RazorAttention addresses this problem since it does not use the attention map as the metric. The head-wise pruning criterion is totally compatible with FlashAttention, and the computation overhead of the compensation token is negligible. Therefore RazorAttention could achieve a substantial inference speedup when compared to previous methods.

To the best of our knowledge, RazorAttention is the first training-free token reduction algorithm that achieves a nearly lossless 3X KV cache reduction. We evaluated RazorAttention on models including Qwen [], Llama-2 [], Llama-3 [] and Baichuan [] on long-context tasks to prove its effectiveness. Our contribution can be summarized as follows:

- •
We systematically analyze the attention dynamic of Transformers under lengthy inputs. Our work reveals that only a few retrieval heads can essentially recall information from the whole input while the rest heads mainly focus on the local context.
- •
We introduce a novel algorithm, namely RazorAttention, that is capable of reducing the KV cache size by $70\%$ under minimal impact on performance for contexts ranging from 8K to 100K tokens. We designed an accurate and data-free metric for allocating all the retrieval heads, together with an error compensation strategy for compensating the information loss due to the truncated KV cache.
- •
RazorAttention introduces negligible overhead in compression and is compatible with FlashAttention, rendering it an efficient and plug-and-play solution that enhances LLM inference efficiency without training or significant overhead. Extensive experiments demonstrate that RazorAttention can be effectively applied to various models and tasks.

## 2 Related Work

As the sequence length increases, the memory consumption of KV cache rapidly expands, potentially surpassing the size of the model parameters themselves. This leads to an urgent need for KV cache compression, particularly in scenarios with limited GPU memory. One direction is non-Transformer architecture design, such as Mamba [], Mamba2 [], Infini-Transformer [], RWKV []and Griffin []. However, in this paper we focus on KV cache reduction for typical Transformers, which is the most widely used model structure. Below we introduce several approaches for KV cache compression.

##### Quantization

Quantization is a classic yet effective approach to neural network compression. In the field of LLM Quantization, while the outlier challenge attracts great attention [, , ] to tackle, the application of which on KV cache is often seen as a by-product of activation quantization. Nevertheless, there are several noteworthy works demonstrating the value of KV cache quantization. FlexGen, Atom and QServe [, , ] carefully designed quantization pipelines that utilize KV cache compression to boost the overall inference throughput. KVQuant [] integrates several techniques to minimize KV quantization error and KIVI [] pushed the limit towards 2-bits. Besides the post-training methods, LLM-QAT [] offers a data-free distillation process that further recovers the performance of the model.

##### Token-dropping

Token-dropping methods assume that not all key-value pairs are essential in self-attention computations, so memory usage can be saved by identifying and removing unimportant KV cache. StreamingLLM [] utilizes sliding window technology, preserving only the KV pairs of attention sink tokens and those within the sliding window, thereby reducing memory footprint and stabilizing model performance. H2O []
is one of the pioneers that use the attention scores to evaluate the importance of each token, followed by an eviction strategy that greedily selects cache with higher scores.
Scissorhands [] and one of the latest work SnapKV [] use similar ideas by narrowing the computation range to consider attention scores related to recent information. Built on that, PyramidKV and PyramidInfer [, ] analyze the attention concentration patterns and further reduce KV cache in later layers. Moreover, research efforts have been made to understand KV cache from different perspectives: FastGen [] paid attention to special tokens and punctuation, SubGen [] investigated the clusterability of key embedding and CORM [] discovered strong correlation amongst tokens of near neighbors.

##### Non-MHA Attention

Another category focuses on reducing KV cache by sharing cache across attention heads. MQA [] aggressively uses a single KV head for all heads, whereas GQA [] suggests an intermediate number of heads to balance the trade-off between inference speed and output quality. Furthermore, MLA [] presents a novel caching method by low-ranking KV cache of all heads into single latent space.

Our algorithm is motivated by the idea from  [], where the authors noticed that there are certain groups of attention heads, denoted as the induction heads, that can effectively recall the queried information from the input. Recent study [] also validated this property under extended inputs. This is the first work that proposes a head-wise pruning criterion for KV cache compression based on the interpretability of the attention mechanism.

## 3 Methodology

In this section, we introduce the key components of RazorAttention. We firstly apply RazorAttention to models using ALiBi [] positional embedding (denoted as ALiBi models) to provide an intuitive understanding of the retrieval and non-retrieval heads. Afterwards, we demonstrate that models using RoPE [] positional embedding (denoted as RoPE models) also exhibit this crucial characteristic, which reveal that KV cache within RoPE models can also be efficiently compressed under minimal loss of accuracy.

### 3.1 RazorAttention for ALiBi models

For ALiBi models, its $h$-th attention head computes the attention score according to

$$ $\displaystyle S_{m\to n}\left(\bm{q};\bm{k}\right)=\bm{q}_{m}\bm{k}_{n}^{ \intercal}-l_{h}(m-n),$ (1) $$

where $\bm{q}_{m}$ is the query tensor at the $m$-th position, $\bm{k}_{n}$ is the key tensor at the $n$-th position, $l_{h}$ is the head-specific slope, $S_{m\to n}\left(\bm{q};\bm{k}\right)$ is the attention score. Notice that $(m\geq n)$ is guaranteed by the casualty of attention.

In the scenario where $l_{h}(m-n)$ significantly dominates $\bm{q}_{m}\bm{k}_{n}^{\intercal}$, the attention between $\bm{q}_{m}$ and $\bm{k}_{n}$ would decay to zero, meaning that the contribution of any tokens positioned further than $n$ becomes negligible for the output at position $m$. The following theorem formalizes this observation.

###### Theorem 1 .

Given an attention head that calculates the attention score as per (1), for any $\epsilon\in(0,1)$, the attention weight from $\bm{q}_{m}$ to $\bm{k}_{n}$ can be upper bounded by:

$$ $\displaystyle\text{Attn}_{m\to n}\left(\bm{q};\bm{k}\right)=\frac{\exp{\left(S _{m\to n}\left(\bm{q};\bm{k}\right)\right)}}{\sum_{n=0}^{m}\exp{\left(S_{m\to n }\left(\bm{q};\bm{k}\right)\right)}}\leq\epsilon,\quad\forall n<m-C_{0},$ $\displaystyle L_{h}:=\frac{2\|W_{Q_{h}}W_{K_{h}}\|_{2}\left(\|\bm{\gamma}\|^{2 }+\|\bm{b}\|^{2}\right)-\log(\epsilon)}{l_{h}}.$ (2) $$

Here $W_{Q_{h}}$ and $W_{K_{h}}$ are the query and key matrices of the $h$-th attention head, $\gamma$ and $\bm{b}$ are the weight and bias for the LayerNorm layer before attention ($\bm{b}=\bm{0}$ for RMSNorm []), and $\|\cdot\|_{2}$ denotes the $l_{2}$-norm of the matrix. $L_{h}$ can be viewed as the vision scope of the head. The detailed proof can be found in Appendix A.

Theorem (1) indicates that when the distance between $\bm{q}_{m}$ and $\bm{k}_{n}$ exceeds $C_{0}$, the attention weight between these two tokens falls below $\epsilon$. When $\epsilon$ is sufficiently small (e.g., 0.1%), remote tokens impose minimal influence on the final output and can thus be discarded. Building on this principle, ALiBi models dynamically adjust the KV cache size for each head. We first compute the effective attention scope $L_{h}$, and keep only the recent $L_{h}$ tokens in the KV cache, since any token further than $L_{h}$ impose attention weight no more than $\epsilon$, we can safely discard them for compression. Therefore, for ALiBi models, the retrieval heads are the ones with a larger $L_{h}$, while the non-retrieval heads has a smaller attention vision $L_{h}$.

###### Theorem 1 .

### 3.2 RazorAttention for RoPE models

For RoPE models, each attention head computes the attention score according to

$$ $\displaystyle S_{m\to n}\left(\bm{q};\bm{k}\right)=\bm{q}_{m}\bm{k}_{n}^{ \intercal},\quad\bm{q}_{m}=\mathcal{R}_{m}\bm{q},\quad\bm{k}_{n}=\mathcal{R}_{ n}\bm{k}$ (3) $$

where $\bm{q}_{m}$ and $\bm{k}_{n}$ are the query and key state after rotary transformation, $\mathcal{R}_{m}$ and $\mathcal{R}_{n}$ are the rotary matrices at position $m$ and $n$ (see [] for details). Although RoPE embedding does not inherently suggest a long-range decaying attention, our empirical findings indicate that the majority of attention heads maintain a limited scope of attention. Notably, only about 15% of the heads, which we term as retrieval heads, are capable of effectively utilizing long-range information while the rest heads only focus on the local context. As shown in Table 1, a significant decrease in accuracy of 16% is observed when the KV cache size is reduced for these retrieval heads. Conversely, dropping remote tokens within non-retrieval heads results in a comparatively minor performance degradation of 1.5%.

**Table 1: We protected the KV cache within different groups of attention heads while keeping only the recent 4K tokens in the rest. Protecting the KV cache in the retrieval heads can retain most of the LLM’s performance, while protecting randomly heads brings no performance gain. This clearly indicates that most of the attention heads only use local context and only retrieval heads can essentially utilize all context information.**
| Protection heads | All | Retrieval heads | Random heads | None |
| --- | --- | --- | --- | --- |
| MultiFieldQA-en | 46.94% | 45.48% | 40.7% | 40.81% |

Based on the findings above, we directly decrease the KV cache for all non-retrieval heads. The performance of the model is mostly retained as shown in Table 1.However, a notable accuracy gap remains, indicating that some information is still being lost. Moreover, the test result on Needle in a Haystack shows a clear performance degradation even when we protect the KV cache of retrieval heads (see our ablation result in Figure 6). To further improve performance, we designed a lightweight and effective way to compress the information in the dropped token into a “compensation token”. The compensation token is defined as

$$ $\displaystyle\hat{\bm{k}}=\frac{1}{N_{d}}\sum_{m\in\{\mathcal{D}\}}\bm{k}_{m}, \quad\hat{\bm{v}}=\frac{1}{N_{d}}\sum_{m\in\{\mathcal{D}\}}\bm{v}_{m}.$ (4) $$

Here $\hat{\bm{k}}$, $\hat{\bm{v}}$ are the compensation tokens for the dropped KV cache, $\{\mathcal{D}\}$ contains the indices of the dropped tokens and $N_{d}$ is the number of the dropped tokens. Afterward, we discard the dropped tokens and augment the KV cache with the compensation token $\hat{\bm{k}}$ and $\hat{\bm{v}}$, where $\{K,V\}$ are the KV cache of the remaining token after rotary transformation. Denoting the compressed KV cache as $\{K,\hat{\bm{k}}\}$ and $\{V,\hat{\bm{v}}\}$, the attention output of the current token follows

$$ $\displaystyle\text{Attn}(\bm{q}_{m},\{K,\hat{\bm{k}}\},\{V,\hat{\bm{v}}\}\})= \frac{N_{d}\exp{\left(\bm{q}_{m}\hat{\bm{k}}^{\intercal}\right)}\hat{\bm{v}}+ \sum_{n\notin\{\mathcal{D}\}}\exp{\left(\bm{q}_{m}\bm{\bm{k}}_{n}^{\intercal} \right)}{\bm{v}}_{n}}{N_{d}\exp{\left(\bm{q}_{m}\hat{\bm{k}}^{\intercal}\right )}+\sum_{n\notin\{\mathcal{D}\}}\exp{\left(\bm{q}_{m}\bm{\bm{k}}_{n}^{ \intercal}\right)}}.$ (5) $$

In Figure 3 we provide an illustrative example of RazorAttention for RoPE models. With compensation tokens, the accuracy is further improved, making RazorAttention almost lossless even dropping $70\%$ of the KV cache in the non-retrieval heads. Below we introduce how we determine the retrieval heads group.

Figure: Figure 3: In Figure 3 we present the illustration of how RazorAttention compress the KV cache. For retrieval heads, we maintain a full cache for retaining all the tokens’ information. For non-retrieval heads, we directly discard remote tokens and compress the discarded tokens into a compensation token whose KV cache is denoted as $\{\hat{\bm{k}},\hat{\bm{v}}\}$. In Figure 3 we provide an illustration example of the echo head and induction head. The current token is “B” and the generated token is “C”. In this case, the echo head would mainly attend to token “B” while the induction head mainly attend to token “C” in previous context.
Refer to caption: x4.png

### 3.3 Identification of Retrieval Heads

For ALiBi models, the attention scope can be directly determined via (2) and KV cache can be dropped accordingly. However, for RoPE models, the retrieval heads need to be identified in a more sophisticated way.
Our investigation reveals that two groups of heads are essential in processing long context, so both of them should be included as retrieval heads as stated below.

- •
Echo head: The head tends to attends back to previous token (referred as echo token) identical to the current token.
- •
Induction head: The head tends to attend to the previous token (namely induction token) that is immediately succeeded by the current token. Basically it attends to the coming token that also exists in previous context.

In Figure 3 we present an illustrative example explaining the echo heads and induction heads. In order to identify the retrieval heads, we generate $K$ (for example, $K=2500$) random tokens, repeat these tokens 4 times, and then use it as the input of the model. This design minimizes semantic dependencies among tokens, thereby allowing a clearer observation of the behavior of echo and induction heads.

Subsequently, we calculated the echo score (attention weight to the echo token) and induction score (attention weight to the induction token) of all words across all heads. The selection of retrieval heads involves the top-$14\%$ attention heads with the highest induction score and top-$1\%$ of attention heads with the highest echo score (see Table 2). Notice that although we only use much fewer echo heads than retrieval heads, our investigation indicates that both heads are crucial for the retrieving performance for LLMs (see Section 4.3 for ablation results).

With the retrieval heads being identified, we hereby introduce RazorAttention for RoPE Models in Algorithm 1.

Figure: Algorithm 1 RazorAttention for RoPE Models

**Table 2: General hyper-parameter settings for experiments in the paper, which leads to 3.125x compression of KV cache under long context input.**
| Hyper-parameter | Settings |
| --- | --- |
| Buffer length | $\max(4000,N/5)$ |
| Induction head protection | top $14\%$ |
| Echo head protection | top $1\%$ |
| Sink token num | 4 |

**Table 3: Performance comparison of RazorAttention and other compression algorithms across various LLMs on LongBench. Notice that the performance of Llama3-8B-Instruct on TREC and LSHT are not applicable (close to 0), hence we do not include their result on Llama3-8B.**
| LLMs |  | NrtvQA | Qasper | MF-en | MF-zh | HotpotQA | 2WikiMQA | Musique | GovReport | QMSum | MultiNews | VCSUM | TREC | TriviaQA | LSHT | Lcc | Average |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Qwen1.5-<br>7B-Chat | All KV | 17.58 | 43.16 | 46.94 | 60.98 | 50.96 | 36.36 | 27.86 | 28.78 | 23.24 | 24.02 | 13.91 | 17.64 | 83.77 | 16.96 | 48.35 | 36.03 |
| StreamingLLM | 6.22 | 24.62 | 18.9 | 34.51 | 20.68 | 12.31 | 5.88 | 3.86 | 3.52 | 20.74 | 3.17 | 8.5 | 36.57 | 13 | 42.51 | 17.00 |  |
| H2O | 16.5 | 38.15 | 40.22 | 51.46 | 50.19 | 35.69 | 27.12 | 28.42 | 22.00 | 22.70 | 14.03 | 18.25 | 83.72 | 16.4 | 47.54 | 34.16 |  |
| RA | 16.63 | 43.1 | 46.66 | 61.08 | 50.49 | 36.1 | 28.79 | 26.68 | 22.59 | 23.96 | 13.83 | 20.87 | 83.83 | 15.66 | 47.85 | 35.87 |  |
| Qwen1.5-<br>72B-Chat | All KV | 28.32 | 46.73 | 48.25 | 63.41 | 55.91 | 46.23 | 34.56 | 32.47 | 22.69 | 24.86 | 15.61 | 71.0 | 91.15 | 46 | 65.05 | 46.15 |
| StreamingLLM | 9.57 | 28.33 | 19.06 | 34.98 | 25.32 | 13.42 | 10.08 | 4.11 | 3.79 | 21.1 | 3.74 | 43.0 | 43.72 | 20.5 | 53.6 | 22.29 |  |
| H2O | 27.98 | 41.45 | 43.69 | 55.93 | 54.77 | 45.16 | 34.61 | 32.24 | 22.35 | 24.36 | 14.5 | 70.0 | 91.15 | 42 | 64.2 | 44.29 |  |
| RA | 27.97 | 46.44 | 47.36 | 63.04 | 55.92 | 46.15 | 34.36 | 32.35 | 22.75 | 24.91 | 15.17 | 71.0 | 91.49 | 46 | 64.68 | 45.97 |  |
| Llama3-8B<br>-Instruct | All KV | 21.84 | 37.04 | 45.07 | 52.34 | 44.63 | 27.28 | 23.04 | 28.18 | 24.54 | 26.26 | 14.41 | 0 | 85.90 | 3 | 30.17 | 35.44 |
| StreamingLLM | 0.61 | 16.29 | 13.41 | 20.05 | 2 | 5.84 | 0.37 | 5.22 | 4.63 | 18.89 | 2.52 | - | 11.54 | - | 26.83 | 9.86 |  |
| H2O | 21.14 | 34.1 | 40.84 | 47.13 | 43.47 | 27.13 | 21.31 | 22.85 | 16.36 | 22.3 | 14.52 | - | 86.17 | - | 30.26 | 32.89 |  |
| RA | 21.16 | 36.22 | 42.88 | 51.93 | 44.07 | 26.89 | 22.03 | 26.56 | 23.86 | 25.83 | 15.69 | - | 85.83 | - | 30.25 | 34.86 |  |
| Baichuan2<br>-13B | All KV | 18.63 | 30.16 | 44.1 | 50.36 | 37.93 | 32.62 | 13.90 | 26.16 | 20.14 | 24.58 | 15.66 | 62.5 | 86.61 | 27.5 | 55.36 | 36.41 |
| StreamingLLM | 5.12 | 12.44 | 23.53 | 32.52 | 16.93 | 16.08 | 6.15 | 5.53 | 1.03 | 5.6 | 3.94 | 42.22 | 30.15 | 7.32 | 35.42 | 16.27 |  |
| H2O | 17.81 | 29.89 | 43.74 | 49.54 | 37.02 | 31.71 | 13.54 | 25.8 | 18.96 | 23.31 | 15.11 | 62.41 | 85.25 | 26.86 | 54.45 | 35.69 |  |
| RA | 18.22 | 31.87 | 43.6 | 51.36 | 36.97 | 32.89 | 13.98 | 25.51 | 20.13 | 24.51 | 15.41 | 62.5 | 87.23 | 28 | 54.53 | 36.45 |  |

Figure: Figure 4: Performance comparison of RazorAttention and other compression algorithms on Llama2-7b-80K, Needle In A Haystack. Notice that H2O is incompatible with FlashAttention so we get OOM errors when tested on longer sequences, and its performance has already become unusable in this case.
Refer to caption: x6.png

## 4 Experiments

A variety of recent-released LLMs are selected to validate our proposals, including Qwen [], Llama2 [], Llama3 [] and Baichuan []. The selected models are evaluated on Longbench [] and Needle In A Haystack [] to demonstrate their capabilities in long-context circumstances. The experiments are conducted on NVIDIA GeForce RTX 4090 (24GB). We will first validate the effectiveness of our proposal on various tasks, followed by the ablation study of each component in our algorithm design. Unless explicitly stated, we use RazorAttention with the hyper-parameters as in Table 2. We use H2O [] and StreamingLLM [] for comparison. Notice that we do not include SnapKV [] as the baseline because it assumes that the query is known before compression, which does not hold in general cases or in a multi-round conversation where the user might query different information from the context (as discussed in Section 1).

### 4.1 LongBench Evaluation

In Table 3 we present the results of different algorithms on LongBench [], which provides a comprehensive assessment to evaluate long-context related abilities of LLMs. We use Qwen1.5-7B and Qwen1.5-72B for testing since they are RoPE models with a context length of 32K. We also include Llama3-8B to validate the performance of RazorAttention on GQA models. We choose Baichuan2-13B to demonstrate the effectiveness of RazorAttention on ALiBi models. It can be seen that RazorAttention achieved a superior performance across all models compared to StreamingLLM and H2O. The compelling outcomes indicate that RazorAttention can achieve comparable performance as the uncompressed baseline, even under 3X compression ratio.

Moreover, we test Llama3-8B-Instruct as a GQA instance where every 4 attention heads share a single set of KV cache. Hence, we consider the attention heads in a group as all retrieval if one or more heads satisfy inductive or echoing property. The results in Table 3 clearly prove that RazorAttention still work for GQA models.

Figure: Figure 5: Adding $1\%$ of the echo heads can significantly enhances the retrieving performance of RazorAttention on Llama2-7B-80k.
Refer to caption: x9.png

### 4.2 Needle In A Haystack Evaluation

In Figure 4 we present the results on Needle In A Haystack. We use Llama2-7B-80K from  [] since the context length of this model is 80K. Unlike H2O, whose performance is severely degraded under long inputs, RazorAttention can still accurately recall the queried information. This is a strong evidence proving that RazorAttention can retain all the semantic information within the original context, while importance-based methods inevitably discard information that might be useful in future queries.

### 4.3 Ablation Studies

Below we present the ablation results of RazorAttention, and prove that the algorithm design and configuration are optimally chosen to achieve a higher compression ration with acceptable performance degradation.

#### 4.3.1 Importance of Echo Heads

Although we only include $1\%$ echo heads in RazorAttention, we notice that this group of heads is quite essential in retrieving information under long context as shown in Figure 5. One possible explanation is that the induction heads depend on the existence of echo heads as discussed in  [].

**Table 4: Qwen1.5-7B-Chat using RazorAttention with different numbers of heads protected, tested on Needle in A Haystack.**
| Protection scheme | Score |
| --- | --- |
| 1% Echo + 5% Induction Head | 69.54% |
| 1% Echo + 8% Induction Head | 78.40% |
| 1% Echo + 11% Induction Head | 84.55% |
| 1% Echo + 14% Induction Head | 86.59% |
| Baseline | 87.05% |

#### 4.3.2 Number of Induction Heads

To determine the optimal number of induction heads to use in RazorAttention, in Table 4 we present the accuracy of RazorAttention under various numbers of induction heads. The results show that the accuracy improves continuously with an increasing number of induction heads. We decide to include $14\%$ of the induction heads in order to achieve an optimal balance between the compression ratio and model performance.

Figure: Figure 6: The compensation token is critical for recovering the information loss introduced by the truncated KV cache.
Refer to caption: x11.png

#### 4.3.3 Importance of the Compensation Token

In Figure 6, it is clearly demonstrated that compensation tokens are critical for the performance of RazorAttention. The compensation tokens successfully compressed most of the information from the dropped tokens,thereby maintaining high accuracy even with significant KV cache reduction.

## 5 Conclusion

In this paper, we propose RazorAttention, a novel KV cache compression algorithm, which successfully achieves a 3X compression ratio for models use RoPE or ALiBi embeddings. Unlike previous importance-based token-dropping methods which inevitably discard semantic information, RazorAttention preserves all semantic information within retrieval heads. We demonstrate that remote tokens can be effectively compressed into compensation tokens within non-retrieval heads. Furthermore, our head-wise pruning criterion is fully compatible with FlashAttention, making RazorAttention a plug-and-play compression method that accelerates the inference of LLMs under extended context. Our experiments demonstrate that RazorAttention can achieve comparable performance with the original model and surpasses previous methods in both accuracy and efficiency.

## 6 Limitation

However, there are still certain limitations of our work. The first question is why attention heads in LLMs behave so differently and how retrieval heads operate under lengthy inputs. The second challenge lies in achieving a higher compression ratio. Although we have successfully reduced the KV cache by $70\%$, we believe this number can be further improved. Moreover, although we have tested our algorithm on several models, the optimal configuration on other models might be different, meaning that we might need more or less retrieval heads under different cases. These topics are quite important and we will keep investigating them in the future work.

## Appendix A Appendix: Proof of Theorem 1

Below we first give an upper bound for the product of the queries and keys, and then show that the attention weight would decay to zero when the positional bias is significantly larger than that upper bound.

###### Proof.

Since we have $\bm{q}=W_{Q_{h}}\bm{x}$ and $\bm{k}=W_{K_{h}}\bm{x}$ where $\bm{x}$ is the input of the Attention block, this leads to

$$ $\displaystyle\bm{q}\bm{k}^{\intercal}=\bm{x}W_{Q_{h}}W_{K_{h}}\bm{x}^{ \intercal}\leq\|W_{Q_{h}}W_{K_{h}}\|_{2}\|\bm{x}\|^{2}.$ (6) $$

Since $\bm{x}$ is attained after LayerNorm, which means

$$ $\displaystyle\bm{x}=\bm{\gamma}\odot\frac{\bm{\hat{}}{x}-\mu}{\sigma}+\bm{b},$ $\displaystyle\mu=\frac{1}{d}\sum_{i=1}^{d}\hat{x}_{i},\quad\sigma=\frac{1}{d} \sum_{i=1}^{d}(\hat{x}_{i}-\mu)^{2}.$ $$

Here $\bm{\hat{}}{x}$ is the input of LayerNorm, $d$ is its dimension and $\hat{x}_{i}$ is the $i$-th dimension of $\bm{\hat{}}{x}$. The equation above leads to

$$ $\displaystyle\|\bm{x}\|^{2}=$ $\displaystyle\left\|\bm{\gamma}\odot\frac{\bm{\hat{}}{x}-\mu}{\sigma}+\bm{b} \right\|^{2}$ $\displaystyle\leq$ $\displaystyle 2\left\|\bm{\gamma}\odot\frac{\bm{\hat{}}{x}-\mu}{\sigma}\right \|^{2}+2\|\bm{b}\|^{2}$ $\displaystyle\leq$ $\displaystyle 2\|\bm{\gamma}\|^{2}+2\|\bm{b}\|^{2}.$ (7) $$

Combining (6) and (7) we get

$$ $\displaystyle\bm{q}\bm{k}^{\intercal}\leq\|W_{Q_{h}}W_{K_{h}}\|_{2}\left(2\| \bm{\gamma}\|^{2}+2\|\bm{b}\|^{2}\right)$ (8) $$

In order to give an upper bound for the attention weight, we have

$$ $\displaystyle\text{Attn}_{m\to n}\left(\bm{q};\bm{k}\right)=$ $\displaystyle\frac{\exp{\left(S_{m\to n}\left(\bm{q};\bm{k}\right)\right)}}{ \sum_{n=0}^{m}\exp{\left(S_{m\to n}\left(\bm{q};\bm{k}\right)\right)}}$ $\displaystyle=$ $\displaystyle\frac{\exp{\left(\bm{q}\bm{k}^{\intercal}-l_{h}(m-n)\right)}}{ \sum_{n=0}^{m}\exp{\left(S_{m\to n}\left(\bm{q};\bm{k}\right)\right)}}$ $\displaystyle\leq$ $\displaystyle\frac{\exp{\left(\bm{q}\bm{k}^{\intercal}-l_{h}(m-n)\right)}}{ \exp{\left(S_{n\to n}\left(\bm{q};\bm{k}\right)\right)}}$ $\displaystyle\leq$ $\displaystyle\frac{\exp{\left(\bm{q}\bm{k}^{\intercal}-l_{h}(m-n)\right)}}{ \exp{\left(\bm{q}\bm{q}^{\intercal}\right)}}$ $\displaystyle\leq$ $\displaystyle\exp{\left(\bm{q}\bm{k}^{\intercal}-l_{h}(m-n)\right)}$ $\displaystyle=$ $\displaystyle\frac{\exp{\left(\bm{q}\bm{k}^{\intercal}\right)}}{\exp{\left(l_{ h}(m-n)\right)}}.$ $$

Therefore to ensure $\text{Attn}_{m\to n}\left(\bm{q};\bm{k}\right)\leq\epsilon$, which is equivalent as $\log\left(\text{Attn}_{m\to n}\left(\bm{q};\bm{k}\right)\right)\leq\log(\epsilon)$, we need

$$ $\displaystyle\log\left(\text{Attn}_{m\to n}\left(\bm{q};\bm{k}\right)\right)\leq$ $\displaystyle\bm{q}\bm{k}^{\intercal}-l_{h}(m-n)\leq\log(\epsilon)$ $$

Taking (8) into the equation above, we get

$$ $\displaystyle\|W_{Q_{h}}W_{K_{h}}\|_{2}\left(2\|\bm{\gamma}\|^{2}+2\|\bm{b}\|^ {2}\right)-l_{h}(m-n)\leq\log(\epsilon),$ $$

which gives us

$$ $\displaystyle m-n\geq\frac{2\|W_{Q_{h}}W_{K_{h}}\|_{2}\left(\|\bm{\gamma}\|^{2 }+\|\bm{b}\|^{2}\right)-\log(\epsilon)}{l_{h}}.$ $$

In this case, we have

$$ $\displaystyle\text{Attn}_{m\to n}\left(\bm{q};\bm{k}\right)\leq\epsilon.$ $$

∎

###### Proof.

## Peer Reviews

confidence: 5
contribution: 3
initial_rating: 8
presentation: 2
questions: DuoAttention and Not All Heads Matter are two concurrent works of RazerAttention. I'd appreciate it if the author could discuss their similarities and differences. In particular, I am interested in Figure 13(1) of DuoAttention — which supposedly features RazorAttention — but the needle performance is very different from your own reporting. I wonder why?
soundness: 3
strengths: 1. The proposed method demonstrates decent performance, providing a significant improvement over many token dropping-based baselines.
2. The evaluation is thorough in terms of model coverages.
3. As recognized by many recent arts, being FlashAttention compatible is a prerequisite of being a pratical long context serving method, and this work fulfills this prerequisite.
summary: This work proposed a method advocating keeping a full KV cache for retrieval/induction heads while adopting a StreamingLLM-like cache pattern for the rests. The proposed method claims to have strong long context performance and many efficiency perks.
weaknesses: 1. The baseline methods (H2O and StreamingLLM) are dated and do not reflect the current SOTA of KV cache compression advancement. I understand if the authors do not intend to feature methods like KIVI or TOVA — as they are with different schools of approaches — but compare with sparse inference techniques like SnapKV or MInference are necessary.

The authors' discussion around L358 regarding SnapKV is faithful, as it corresponds to the findings of SharedContextBench (also submitted to ICLR here). However, SnapKV — and, in essence, methods like LoCoCo and LESS — are still worth featuring since we can almost always find weaknesses in any method. The authors should consider adding some quantifiable evaluations to reflect the weakness of SnapKV-like methods, if they intend to highlight the superiority of RazerAttention in this regard.

2. The needle setting is unclear. It is often very easy to do the "eating sandwich in a park" needle with the "grass sky sun" repetition background. But it is hard with the "magic city number" needle with a noisy background (e.g., Paul Graham's essay). The authors should clarify this setting and consider adding more standardized needle evaluations, like the RULER.

3. Similar to #2, the LongBench truncation settings of Table 3 can also use some clarifications.

4. RazorAttention is proposed as an efficiency work. However, the efficiency evaluation is very limited. I'd like to see latency and throughput evaluation under different context length/batch size workloads.

5. The core design of the work heavily relies on the findings of Induction Heads and Retrieval Heads, making its novelty limited. While I don't think this is much of an issue due to the practical performance gain the proposed method achieved, the relationship of such works deserves much better highlight. The mere mentioning around L167 is not enough.

6. I'd like to see more dataset coverage beyond LongBench — as it is no longer very long w.r.t. models like Llama 3.1. Evaluation of Llama 3.1 on $\infty$Bench would fill this gap.


In summary, I think this is a decent work that cleverly leveraged some observations of retrieval/induction ideas. I am open to increasing my rating should the author provide the requested clarifications and evaluations. I also apologize for the slightly delayed review — I will try to be more engaging during the discussion period.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: See the weakness
soundness: 3
strengths: The paper observes the phenomenon of "retrieval heads" aligning with observations from other works.

Basing on the retrieval heads, the paper presents a straightforward token dropping based KV compression mechanism.

Experiments results show improvement compare with the previous token dropping based baselines.
summary: The paper proposed a novel KV cache mechanism basing on retrieval heads. To be more specific, the paper presents the observation of some retrieval heads that is essential for handling long context workload, then proposed a token-dropping based method to compress the to KV cache of the non-retrieval heads with compensation tokens, while keeping the retrieval head integrated, leading to substantial accuracy improvement compare to previous proposed method.
weaknesses: Lack of experimental results on compression ratio:
The paper compares against StreamLLM and H2O. However, StreamLLM has a very limited compression ratio, and the performance of different compression ratio in H2O also will be different. The current manuscript lack the study of how different compression ratio in RazorAttention will affect the performance.

Lack of overhead evaluation on efficiency
The paper claims that RazorAttention can enhance LLM inference efficiency without overhead. However, there is no efficiency evaluation (both peak memory usage as well as the inference latency analysis). It is essential to have such real efficiency gain evaluation.

The novelty of "retrieval heads"
Previous paper [1] has demonstrated the existence of "retrieval heads" and its importance in serving long context. Do the authors aware of this and is there a difference between the two "retrieval heads" proposed in the two paper?

Miss of clarification:
"Retrieve and Process" in line 84: The manuscript claims that some heads are used for retrieval while some are used for process. However, lacks the analysis of "process"
"Protect" in line 234: what is the specific protection mechanism? 

Potential typos:
115 as follows:

[1] Retrieval Head Mechanistically Explains Long-Context Factuality

---

confidence: 5
contribution: 2
initial_rating: 6
presentation: 3
questions: Please address the weaknesses above.
soundness: 3
strengths: 1. The observation is neat.  And experiments support the hypothesis quite well. 
2. The recipe to identify the heads is very lightweight (although is this a contribution or is it previously known from anthropic paper?)
summary: The paper has made an observation: induction and echo heads are retrieval heads and others are local processing heads. They then keep complete KV Cache for retrieval heads and drop non-local KV for local processing heads. Following their hyper parameters, this leads to 70% asymptotic KV Compression.
weaknesses: 1. Experiments

     A. why are some datasets from longbench missing like passage_retrieval etc

     B. Can you perform experiments on infinity benchmark for comparison.

     C. what is the exact KV compression in the longbench datasets? The 70% number is misleading since longbench generally has small context lengths and authors have used a context thold of 4000.  Please add exact KV Cache compression numbers for all the methods to the table. 

2. Novelty: Can you please elaborate on the novelty of defining / identifying the echo and induction heads. what is already known and what is figured out?

3. Writing errors: Please revise the manuscript for text/language errors. e.g. : 
    
     A. "There are been plenty ..." (page 1)  
    
     B. " and proved that accuracy .." ( page 2) {you showed it empirically .. and also 70% is not what you show in the experiments.}.

     C. C_0 in equation 2. Do you mean L_h ?

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: - What are the latency and throughput improvements, both absolute and relative to the baseline?
- What are the comparative performances at different KV cache compression ratios?
- Has RazorAttention been tested on other models like Mistral-7B-Instruct-v0.2 or longchat-7b-v1.5-32k?
- Can you provide a breakdown or approximation of the distribution between retrieval and non-retrieval heads?
soundness: 3
strengths: - RazorAttention’s head-specific strategy, distinguishing retrieval heads from non-retrieval heads, is a novel approach to KV cache compression that aligns with the attention heads’ different roles.
- This method is compatible with FlashAttention and offers practical efficiency gains without training or model alterations.
- RazorAttention is tested across various architectures, including ALiBi and RoPE models, proving effective on different LLM types, with compression ratios of up to 70%.
summary: The paper introduces RazorAttention, a training-free algorithm for compressing KV cach in LLMs. This method selectively retains information in specific “retrieval heads” of the multi-head attention mechanism while discarding remote tokens in other heads. A notable aspect of RazorAttention is using a "compensation token" to summarize discarded information, minimizing data loss. RazorAttention shows a 70% reduction in KV cache size with minimal performance impact, supporting compatibility with FlashAttention for efficient large-context inference without modifying the base model.
weaknesses: - The paper lacks detailed efficiency reports, particularly on latency and throughput improvements. A breakdown of the actual performance gains would give better insight into its practical benefits.
- RazorAttention introduces operations like compensation token creation and slicing, yet the associated computational overhead is not profiled and discussed. A detailed overhead analysis and implementation would help.
- RazorAttention’s comparison is primarily with StreamingLLM and H2O, omitting comparison and discussion of other types of KV cache compression methods. (like what are others’ performances under similar compression rates? can this method be combined with others?)
- The results focus on one compression rate without experiments on varying compression ratios compared with others.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: 1.	What is the compression ratio for GQA in Llama3-8B-Instruct model? What is the compression ratio for 1024K sequence length?
2.	How to decide the 1% echo heads and 14% induction heads configuration? Have the authors tried different combination, e.g. more echo heads and less induction heads?
soundness: 3
strengths: 1.	It is pluggable solution that enhances the inference efficiency of a diverse set of large language models without noticeable impacts on performance.
summary: This paper introduce a KV cache compression algorithm in LLMs that only maintains a full cache for retrieval heads and reduces all the remote tokens in non-retrieval heads to single compensation token per head. The proposed method can achieve 3X KV cache compression ratio without noticeable impacts on performance.
weaknesses: 1.	There are some related works identifying the different patterns of retrieval and non-retrieval heads, such as DuoAttention, Retrieval Head Mechanistically Explains Long-Context Factuality, MInference, etc. It would be better to clarify the novelty of the proposed method.
2.	It can only be applied to multiple KV cache scenarios. For only single KV cache model structure, such as YOCO, it cannot be applied.
3.	The baselines are not comprehensive. It would be beneficial if the authors could compare their results with more advanced baselines.
4.	It seems the proposed method cannot achieve the best performance in all the benchmarks. It would be beneficial to give detailed analysis of the worse cases.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. As mentioned in the weaknesses, it is beneficial to compare with DuoAttention in terms of effectiveness.
2. As a training-free method, a potential issue is that different models may have different optimal hyperparameters. For example, how can we quickly determine the suitable proportion of Induction Heads and Echo Heads for different models?
3. If a general theoretical analysis can be found explaining why retrieval heads are widely present in various models, that would be highly significant.
soundness: 3
strengths: 1. The research problem addressed in this paper is significant. Reducing the key-value  cache in large language models is a hot topic, and improvements in this area can have a substantial impact on the real-world deployment of LLMs.
2. Based on the phenomenon of retrieval heads present in large language models and in conjunction with the token dropping method, this paper explores an training-free way to reduce the key-value cache. The results from downstream tasks and NIH tests have both demonstrated the effectiveness of this approach.
3. The writing of this article is easy for readers to understand.
summary: The paper introduces RAZORATTENTION, motivated by the goal of reducing the size of the key-value cache as much as possible without losing semantic information. The rationale behind this approach is the retrieval heads in LLMs. For these retrieval heads, RAZORATTENTION retains the KV cache in its original state; for the other heads, it only caches recent tokens and attention sinks. Additionally, the paper designs a compensation strategy for compensating the information loss. Experiments indicate that RazorAttention can compress up to 70% of the KV cache without causing a noticeable decline in performance.
weaknesses: 1. I agree that this work is the first training-free token reduction algorithm based on retrieval heads. However, to my knowledge, the phenomenon of retrieval heads being prevalent in LLMs was first articulated in [1]. [1] provides a detailed discussion on the properties of retrieval heads, which diminishes the novelty of this paper. Although the paper provides a theoretical analysis of AliBi-based models, which is commendable, it does not cover Rope-based models.
2. A recent work [2] is also motivated by the presence of retrieval heads. The main idea of [2] is to combine streaming heads with retrieval heads, which is very similar to the approach taken in this paper. But [2] is not training-free. It would be beneficial to elaborate on the advantages of this paper over [2]. 

[1] Retrieval head mechanistically explains long-context factuality.
[2] DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: The paper lacks detailed experimental setup descriptions. When comparing with baseline methods, was the total kv budget (i.e., the combined kv cache size for all heads) kept constant?
soundness: 2
strengths: The paper is well-written and straightforward to follow.
summary: This paper investigates the attention mechanisms within transformer models, observing that most attention heads primarily focus on local context while only a few, termed "retrieval heads," are capable of attending to all input tokens. To optimize memory and computational efficiency, the authors propose maintaining a full key-value (kv) cache for retrieval heads and only keeping a recent kv cache for other attention heads. The identification of these retrieval heads is based on observed patterns known as the "Echo head" and "Induction head."
weaknesses: 1. The authors identify retrieval heads as echo heads and induction heads, demonstrated through an example using a synthetic repeated sequence. However, this example appears specifically designed to highlight these two types of heads. It raises the question of whether this pattern holds in more generalized query settings. Intuitively, it does not seem that in every type of query, the most crucial information is located solely in echo or induction tokens. Would these patterns still be dominant in more diverse and complex input data?

2. The authors state that they did not compare their method with SnapKV because SnapKV relies on the query for selecting the kv cache. However, this reasoning is weak, as SnapKV is a well-known and accepted method in the field. A comparison would provide more comprehensive insights.

3. I'm confusing that whether the selection of retrieval heads is model-based or  also query-specific. Are these retrieval heads consistent across different queries, or do different heads have different functions across different queries?

4. The authors note that reducing the kv cache for all non-retrieval heads introduces a notable accuracy gap and that compensation tokens can mitigate this gap. However, it is unclear whether the observed improvement comes purely from the use of compensation tokens or from the strategic allocation of the kv cache itself. An experiment comparing the combination of compensation tokens with other baseline models could clarify the extent of the improvement attributable to each factor.

Meta-review: The paper introduces RazorAttention, a novel KV cache compression algorithm that selectively retains full information in retrieval heads while condensing remote tokens in non-retrieval heads using a compensation token. This method achieves over 70% reduction in KV cache size without compromising model performance, offering compatibility with FlashAttention and requiring no retraining. The key strengths include its efficient head-specific strategy, lossless compression, and extensive evaluation across diverse language models, showing practical utility and innovation. However, limitations include the reliance on earlier insights about retrieval heads, limited comparisons with some advanced baselines, and a need for more detailed efficiency evaluations on GPUs. Despite these, the paper demonstrates strong practical performance and a significant impact on LLM efficiency, making it a valuable contribution to the field. The decision to accept is based on its substantial improvements in resource efficiency, innovative methods for KV cache compression, and broad applicability.

Decision: Accept (Poster)