In-context KV-Cache Eviction for LLMs via Attention-Gate

## Abstract

Abstract The KV-Cache technique has become the standard for the inference of large language models (LLMs).
It caches states of self-attention to avoid recomputation.
Yet, it is widely criticized that KV-Cache can become a bottleneck of the LLM inference system, especially when confronted with ultra-large models and long-context queries.
A natural remedy is to discard the KV-Cache for less important tokens, with StreamingLLM (Xiao et al., 2024 ) as an example, but the used static eviction strategies cannot flexibly adapt to varying contexts.
Remedies like H2O (Zhang et al., 2024 ) leverage accumulative attention scores to perform dynamic eviction but suffer from the attention bias issue in capturing contextual information.
This paper bridges this gap by devising a parameterized KV-Cache eviction mechanism, dubbed as Attention-Gate , which accepts the whole context as input and yields eviction flags for each token to realize in-context eviction.
The subsequent self-attention module proceeds according to the flags and only the KV states for the remaining tokens need to be cached.
The Attention-Gates can vary among different heads and layers and be trivially plugged into pre-trained LLMs, tuned by cost-effective continual pre-training or supervised fine-tuning objectives to acquire what to discard.
The computational and memory overhead introduced by Attention-Gates is minimal.
Our method is validated across multiple tasks, demonstrating both efficiency and adaptability.
After a highly efficient continual pre-training, it achieves higher average accuracy and evicts more tokens compared to traditional training-free methods.
In supervised fine-tuning, it not only evicts many tokens but also outperforms LoRA-finetuned LLMs on some datasets, such as RTE, where it improves accuracy by 13.9% while evicting 62.8% of tokens, showing that effective eviction of redundant tokens can even enhance performance.

## 1 Introduction

Large language models (LLMs) (Dubey et al., ; Team et al., ; Chiang et al., ) have achieved remarkable success across a wide range of natural language processing tasks.
A key technique that has enabled efficient LLM inference is KV-Cache, which stores transient attention keys and values to avoid recomputation.
However, as the size of LLMs continues to increase and the demand for handling long-context queries grows, the KV-Cache has emerged as a significant bottleneck.
Storing attention states for numerous tokens can lead to considerable memory overhead and data transfer among the memory hierarchy results in substantially increased inference time.

Studies have shown that sparsity is a natural phenomenon in attention mechanisms, with many tokens being redundant for inference (Zhang et al., ).
This suggests that retaining all tokens in the KV-Cache is unnecessary. Existing works have explored this insight
to compress KV-Cache using static strategies or hinging on accumulative attention scores.
StreamingLLM (Xiao et al., ) is a representative of the former by retaining a fixed window of beginning and recent tokens in the KV-Cache but it struggles to flexibly adapt to specific contexts.
E.g., in sentiment analysis, retaining the token “cute” in “a cute cat” is crucial, while in object recognition, the token “cat” would be more important.
H2O (Zhang et al., ), on the other hand, employs a token-adaptive approach, using local accumulative attention scores to determine which tokens to evict.
However, it is criticized that in practice, H2O suffers from the attention bias issue (Oren et al., ), with a tendency to over-prioritize either the initial or recent tokens.

Figure: Figure 1: KV-Cache eviction patterns across different layers and attention-heads, visualized for 4 samples from the PIQA dataset (top row) and 4 samples from the BoolQ dataset (bottom row), using AG fine-tuned Llama2-7B models. Black areas represent tokens that are neither computed nor stored in the KV-Cache. The variability of eviction patterns across tasks, prompts, layers, and attention-heads demonstrates the dynamic nature of our method. A common trend observed is that deeper layers tend to mask more KV-Cache states, with some in deeper layers being entirely masked.
Refer to caption: x1.png

To overcome these challenges, we introduce a parameterized KV-Cache eviction mechanism named the Attention-Gate (AG), to perform reliable *in-context* eviction.
AG is positioned before a self-attention layer within the model. It takes a sequence of token features as input and generates eviction flags for the tokens, indicating whether a token should be excluded from subsequent self-attention computations.
Tokens that are evicted do not require their KV states to be cached.
AGs can be seamlessly integrated into pre-trained LLMs and tuned by minimizing the language modeling loss.
Ideally, AGs can automatically learn to discern the most relevant tokens for the current context without manual intervention.
In practice, we can implement the AG as a self-attention layer with much fewer heads than the original model (e.g., 4 v.s. 32).
This way, the parallel computational capabilities of the hardware can be harnessed to minimize the extra overhead introduced by AGs.

AG is empirically shown to enjoy high training efficiency, e.g., only four NVIDIA 4090 GPUs and a dataset of 5,000 samples are required for continual pre-training when applying AGs to LLaMA2-7B (Touvron et al., ).
This alleviates concerns about the computational overhead related to trainable eviction strategies (Zhang et al., ; Chen et al., ) and amplifies the performance superiority of our approach over existing training-free approaches.
As illustrated in Figure 1, AG generates different eviction strategies across different layers and attention-heads for different tokens, demonstrating its adaptability to the diverse requirements of each component in the model.

To validate the effectiveness of our method, we conduct extensive experiments across multiple benchmarks. After efficient continual pre-training (CPT), our approach outperforms traditional training-free eviction strategies, such as StreamingLLM and H2O, in both accuracy and token eviction rates.
In supervised fine-tuning (SFT), our method not only evicts a significant number of redundant tokens but also maintains or surpasses the performance of LoRA-finetuned LLMs. For example, on the RTE dataset, our approach improves accuracy by 13.9% while evicting 62.8% of tokens, demonstrating that selective token eviction can enhance performance.
In summary, the Attention-Gate mechanism provides a scalable and efficient solution for KV-Cache management, addressing the limitations of traditional training-free methods.

Figure: Figure 2: An overview of Attention-Gate (AG) for KV-Cache eviction. AG is a lightweight learnable module placed before each MHA layer. Given the input hidden states, it determines for each head whether to retain or discard the key and value tokens in the KV-Cache. In the attention weights, this corresponds to masking out columns for the evicted keys, while keeping the diagonal intact to ensure the query interacts with its own key.
Refer to caption: x2.png

## 2 Related Work

As large language models (LLMs) scale in both size and input sequence length, optimizing their efficiency has become increasingly important, particularly in addressing space and time complexity. A significant bottleneck lies in the attention mechanism, which demands considerable computational and memory resources, especially for long sequences.

Computation and memory challenges.
The computational cost of attention is driven by its quadratic complexity, while the memory burden stems from the need to store the KV-Cache during inference to avoid recomputation. Although several works have focused on reducing the computational load of attention mechanisms, many fail to address the memory constraints that remain a critical issue. Techniques like Reformer (Kitaev et al., ) and FlashAttention (Dao et al., ; Dao, ) mitigate memory usage for long sequences but still require substantial cache storage. On the other hand, approaches such as MQA (Shazeer, ) and GQA (Ainslie et al., ) prioritize memory compression without tackling computational complexity. These solutions tend to focus on either memory or computational efficiency, rarely addressing both simultaneously.

KV-Cache eviction strategies.
To address both memory and computational challenges, KV-Cache eviction has emerged as an effective strategy. Existing approaches can be categorized into static and adaptive methods based on accumulative attention scores.

Static strategies, such as those used in Sparse Transformers (Child et al., ), employ fixed pruning patterns, such as Strided and Fixed Attention. While effective in some cases, these approaches are not adaptive to specific contexts, often sacrificing accuracy. StreamingLLM (Xiao et al., ) tackles the *Attention Sink* phenomenon, where attention scores concentrate on initial tokens, by retaining these tokens along with a fixed window of recent tokens. While this improves performance, static approaches generally lack the flexibility needed to adapt to different tokens, attention-heads, or layers.

Strategies using accumulative attention scores offer more flexibility by dynamically identifying important tokens. For instance, SpAtten (Wang et al., ) employs Accumulative Attention Scores (A2S), which sum the softmax outputs for each token to measure its importance. This approach allows selective token pruning in subsequent layers, effectively reducing computational complexity without the need for retraining. H2O (Zhang et al., ) extends this concept to decoder-based models, using local A2S statistics for adaptive eviction in autoregressive generation. However, H2O suffers from the attention bias issue (Oren et al., ), particularly in long-context inputs.
Several follow-up works have aimed to address this limitation. NACL (Chen et al., ) introduces random eviction to mitigate attention bias, while A2SF (Jo & Shin, ) incorporates a Forgetting Factor. However, none of these approaches fully resolve the underlying problem.

More Adaptive strategies.
Although strategies based on accumulative attention scores provide more flexibility than static methods, they still have notable limitations. For instance, H2O (Zhang et al., ) applies the same token eviction ratio across all attention heads, restricting the adaptability of the method.
FastGen (Ge et al., ), on the other hand, introduces a different approach by hybridizing KV-Cache compression policies and applying adaptive strategies to each attention head. However, it focuses on the decoding stage and neglects the importance of the prefilling stage.
Learnable eviction strategies, on the other hand, offer greater flexibility by enabling different layers and attention heads to adopt heterogeneous eviction policies. However, such strategies have been relatively underexplored, likely due to concerns about the computational overhead they may introduce (Zhang et al., ; Chen et al., ). Nonetheless, task-specific training is essential for optimizing performance across different contexts. For example, a recent approach (Anagnostidis et al., ) introduces a learnable mechanism for dropping uninformative tokens, but it faces difficulties in batched generation and does not account for continual pre-training or decoding-only LLMs.
Despite these challenges, learnable strategies have strong potential to improve performance across a variety of tasks by allowing models to adapt their eviction strategies to meet task-specific requirements.

## 3 Method

This section first briefly reveals the basics of multi-head attention and KV-Cache and then describes the proposed *Attention-Gate* (AG) mechanism for in-context KV-Cache eviction for LLM inference acceleration.
An illustrative overview of AG is presented in Figure 2.

### 3.1 Preliminary

Multi-Head Attention (MHA) (Vaswani et al., ) is a core component of the Transformer architecture, widely used in most LLMs.
MHA enables the model to capture dependencies across different tokens in a sequence by computing self-attention across multiple attention-heads.
Each head attends to different parts of the input sequence independently, allowing the model to capture various aspects of the relationships between tokens.
The outputs of these attention heads are then concatenated and projected through a final output matrix.

KV-Cache is employed to store the key and value representations of tokens from previous time steps during inference.
This prevents redundant recomputation of these representations for every new token generated, significantly speeding up auto-regressive generation.
The inference process in auto-regressive transformers, such as Llama (Touvron et al., ), can be divided into two stages: *prefilling* and *decoding*.
1.
In the prefilling stage, the model processes the entire input sequence, generating key-value pairs for each attention-head in each layer.
These key-value pairs are stored in the KV-Cache for future reuse.
2.
In the decoding stage, when generating the next token, the model only computes the key and value for that new token and appends them to the cached key-value pairs.
This reduces the computational load, as the previously cached key-value pairs do not need to be recomputed.
The process continues token by token until the sequence generation is complete.

KV-Cache plays a critical role in improving the efficiency of LLM inference, especially in scenarios where long sequences are processed. However, the size of the KV-Cache grows with the input sequence length, leading to substantial memory overhead. Efficiently managing this cache while maintaining model performance has become a key challenge in scaling LLMs to longer contexts.

### 3.2 Limitations of Traditional Eviction Strategies

Lack of flexibility.
Flexibility in KV-Cache eviction strategies spans across various dimensions, including token-specific, attention-head-specific, layer-specific, task-specific, and model-specific adaptability.
Static eviction policies offer no such flexibility, requiring manual intervention to adjust them.
While methods based on accumulative attention scores (e.g., H2O (Zhang et al., )) have improved upon static strategies by introducing token-level and head-level adaptability, they are still limited.
Specifically, in H2O, the eviction ratio remains uniform across all attention heads, restricting head-level flexibility.
Flexibility is crucial in eviction strategies because models process varying types of data across different contexts, requiring the ability to selectively retain important tokens while discarding others. Without this adaptability, the model may inefficiently store redundant information, increasing memory usage and slowing down inference times.

Absence of global statistics.
Eviction strategies that rely on accumulative attention scores, such as H2O, base their calculations on local statistics, which often results in attention bias issues (Oren et al., ).
This local approach may misjudge a token’s importance, especially when the context spans long sequences.
While several methods (Chen et al., ; Jo & Shin, ) have proposed solutions to mitigate this bias, they are still fundamentally rooted in local statistics, addressing the issue only partially.
A more effective solution would incorporate global statistics to ensure that the eviction process is based on a more comprehensive understanding of the context, thereby reducing biases and improving the accuracy of token retention.

Inefficiency.
Efficiency is another challenge for traditional methods such as H2O and FastGen (Ge et al., ). These methods operates in a token-by-token manner for each decoding step,
sequentially deciding which key-value pairs to evict from the cache. Additionally, H2O determines which tokens to evict after computing the attention scores, meaning some computation is wasted on tokens that are later discarded. This sequential approach hinders efficiency, especially when handling large batches of data or when processing long sequences.

Our Attention-Gate mechanism is designed to address the aforementioned limitations.

### 3.3 Attention-Gate

The *Attention-Gate* (AG) is a lightweight, trainable module designed to determine which tokens in the KV-Cache of each attention-head should be retained or discardedAG is positioned before the Multi-Head Attention (MHA) layer and operates by generating *eviction flags* for each token, guiding both the computation of attention scores and the management of the KV-Cache.

Input and output.
AG takes the hidden states as input and generates binary flags as output, one for each token in every attention head.
These flags indicate whether the corresponding token’s key-value pairs should be stored or discarded in the KV-Cache.
Specifically, the input is a hidden state matrix $X\in\mathbb{R}^{n\times d}$, where $n$ represents the sequence length, and $d$ denotes the dimensionality of the hidden states.
The output consists of independent binary decisions for the retention or eviction of each token’s key-value pairs across all attention heads in the MHA module.

Local or global?
The structure of AG offers two choices: using either local or global information to guide the eviction policy.
1.
Local information: The simplest approach leverages a *linear layer* where each token only considers its own hidden state without taking other tokens in the sequence into account.
This method is computationally efficient and straightforward, making it appealing for its simplicity.
2.
Global information: On the other hand, a more comprehensive approach adopts an attention-like structure, where tokens can aggregate information from across the entire sequence.
This enables the eviction policy to reflect the global in-context information, making decisions that are informed by the broader context.
Naturally, this method is computationally heavier than the local approach, but it offers a more accurate reflection of the sequence’s token importance.

In our design, we opted for the attention-like structure because *global context is essential for making effective eviction decisions*.
As demonstrated in Section 4.4, using an attention-like structure significantly outperforms the linear layer, as the latter struggles to properly evict tokens.
This outcome is expected, as determining whether a token is redundant or crucial for inference often requires a global understanding of the sequence.
A local approach, like a linear layer focusing on individual tokens, lacks the broader perspective necessary to make accurate decisions about which tokens are important.
This finding also underscores the limitations of methods like H2O, which rely on local statistics rather than global ones, as mentioned in Section 3.2.

Softmax or Sigmoid?
After gathering information through the attention-like structure, the next step is to compute the eviction probability for each token. Common activation functions such as *Softmax* and *Sigmoid* can be used to generate these probabilities.

In our case, we choose the Sigmoid function for a key reason: while Softmax normalizes probabilities across the entire sequence, introducing competition between tokens, Sigmoid treats each token independently.
This independence is crucial because we do not want a token’s likelihood of being evicted to depend on the eviction probabilities of other tokens in the sequence.

Eviction flags.
After computing the eviction probabilities for each token, we apply a threshold $\tau$.
If the probability of retaining a token exceeds $\tau$, that token is retained in the KV-Cache; otherwise, it is discarded.
In the attention matrix, the columns corresponding to evicted tokens are masked out.
However, the diagonal elements, where a token attends to itself, are always preserved.
This guarantees that each token continues to interact with its own key.

Prefilling or decoding?
AG is primarily applied during the prefilling stage, where the full sequence is available.
By making eviction decisions before the MHA layers, AG effectively manages the KV-Cache during this phase.
In contrast, AG is not used during the decoding stage to avoid adding complexity to the inference process.
Using AG in decoding could slow down inference and increase the cache memory footprint, as additional keys and values from the attention-like structure would need to be stored.

In summary, AG dynamically determines which tokens should be retained, optimizing both KV-Cache usage and attention score computation by utilizing an attention-like structure, Sigmoid activation function, and a threshold-based decision for token eviction.

### 3.4 Training Implementation

To train the Attention-Gate (AG) effectively, this section outlines the key components of the training process. For more comprehensive details, please refer to Section 4 and Appendix A.

Eviction Loss.
To encourage the eviction of unnecessary tokens, we introduce a dedicated loss function called the *Eviction Loss*.
This loss encourages the model to discard as many tokens as possible, which is defined as:

$$ $\ell_{\text{evict}}=\alpha\cdot\frac{1}{L}\sum_{l=1}^{L}\text{ReLU}\left( \overline{\text{AG}}_{l}-\beta\right),$ (1) $$

where $L$ is the number of layers, and $\overline{\text{AG}}_{l}$ represents the average output of the AG module in layer $l$.
In this formula, $\alpha$ adjusts the intensity of KV-Cache eviction, while $\beta\in[0,\tau]$ ensures that eviction does not become overly aggressive.This loss function works alongside the auto-regressive loss to balance token eviction with maintaining model performance.
The loss allows for adaptive eviction across layers and attention-heads. As shown in Figure 1, after training, the model learns to increase token eviction in deeper layers, which is not achievable with training-free methods like H2O (Zhang et al., ) or StreamingLLM (Xiao et al., ).

Initialization.
We initialize the AG parameters using Xavier initialization (Glorot & Bengio, ) to provide a stable starting point for learning.
Additionally, a small constant $\gamma\geq 0$ is added inside Sigmoid, ensuring that the initial retention probabilities are close to 1.
This encourages the model to retain most tokens early in training, allowing it to learn which tokens are important more gradually.

Handling non-differentiability.
Directly applying the threshold-based gating mechanism from Section 3.3 would lead to non-differentiable gradients during training due to the hard thresholding’s discrete nature.
To resolve this, we employ the *Straight-Through Estimator* (STE) (Yin et al., ), which allows gradients to flow through discrete decisions by approximating them during the backward pass.
Specifically, during backpropagation, instead of using the hard 0 or 1 values obtained from comparing against the threshold, we utilize the smooth output of the Sigmoid function.
This approach ensures smooth gradients and enables effective training of the AG while preserving its binary behavior during the forward pass.

## 4 Experiments

This section consists of three main parts.
First, we evaluate the performance of AG in two scenarios: continual pre-training (CPT) and supervised fine-tuning (SFT).
For CPT, we compare AG with classic methods on performance and eviction rates (Table 2).
For SFT, we benchmark AG against vanilla LoRA fine-tuning to highlight its task-specific adaptability (Section 4.2).
Second, we provide visualization of selected examples to demonstrate the core characteristics of AG (Section 4.3).
Finally, we conduct ablation studies to provide further insights into the effectiveness of AG (Section 4.4).

### 4.1 Continual Pre-training

**Table 1: Performance comparison of Llama2-7B and various KV-Cache eviction strategies across seven tasks. The table reports accuracy (Acc.) for Llama2-7B and all eviction methods, with Llama2-7B serving as the upper bound for accuracy. Metric %Eviction refers to the mean KV-Cache eviction ratio, representing the percentage of tokens evicted from the KV-Cache. The eviction ratio is fixed at 50% for the baseline methods, including a local strategy (retaining only recent tokens), StreamingLLM, and H2O. In contrast, our method, which includes continual pre-training, achieves higher average accuracy while maintaining a higher average %Eviction.**
|  | Metric | PIQA | ARC-C | ARC-E | RTE | COPA | BoolQ | HellaSwag | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Llama2-7B | Acc. | 76.33 | 37.29 | 51.32 | 51.99 | 62.00 | 69.94 | 68.16 | 59.58 |
| Local | Acc. | 69.97 | 31.86 | 48.68 | 51.99 | 60.00 | 57.86 | 37.08 | 51.06 |
| StreamingLLM | Acc. | 72.69 | 33.22 | 51.15 | 50.18 | 63.00 | 62.05 | 40.85 | 53.31 |
| H2O | Acc. | 75.9 | 33.22 | 52.03 | 52.71 | 47.00 | 67.37 | 66.32 | 56.36 |
| Ours | Acc. | 76.17 | 33.90 | 49.03 | 52.35 | 63.00 | 67.52 | 66.33 | 58.33 |
| %Eviction | 54.29 | 51.03 | 51.05 | 46.70 | 40.02 | 57.75 | 52.16 | 51.87 |  |

**Table 2: Performance of Llama2-7B with LoRA fine-tuning and our method on six downstream tasks. In addition to the LoRA fine-tuning targets, our method makes the AG modules learnable. Two settings for $\alpha$ (0.5 and 1) are tested. Our method maintains comparable or better accuracy while achieving a higher eviction ratio, demonstrating its task-specific adaptability in managing token eviction without significant accuracy loss.**
|  | Metric | PIQA | ARC-C | RTE | COPA | BoolQ | OBQA | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Fine-tuned Llama2-7B | Acc. | 82.92 | 60.34 | 64.98 | 92.00 | 88.10 | 78.80 | 77.86 |
| Ours $(\alpha=1)$ | Acc. | 82.15 | 59.66 | 64.26 | 93.00 | 86.82 | 78.80 | 77.45 |
| %Eviction | 66.16 | 48.31 | 65.47 | 45.40 | 67.46 | 67.17 | 60.00 |  |
| Ours $(\alpha=0.5)$ | Acc. | 81.50 | 57.63 | 74.01 | 95.00 | 87.00 | 79.20 | 79.06 |
| %Eviction | 64.96 | 36.45 | 62.80 | 34.77 | 67.31 | 66.65 | 55.49 |  |

#### 4.1.1 Setup

Models and datasets.
We use Llama2-7B (Touvron et al., ) as our primary base model due to its strong performance and popularity within the AI community.
Additionally, we validate the feasibility of our approach on Mistral-7B (Jiang et al., ), with results provided in Table 4.
For continual pre-training, we select a subset of the Redpajama dataset (Computer, ) containing approximately 5,000 samples (^1^11We sampled 4,997 samples proportionally from each subset of RedPajama.).
To assess the effectiveness of our method, we evaluate it on seven widely recognized benchmarks: PIQA (Bisk et al., ), ARC-C (Clark et al., ), ARC-E (Clark et al., ), RTE (Bar-Haim et al., ), COPA (Roemmele et al., ), BoolQ (Clark et al., ), and HellaSwag (Zellers et al., ).
All evaluations are conducted in a zero-shot setting, with accuracy assessed using OpenCompass (Contributors, ).

Training details.
During continual pre-training (CPT), only the AG modules are learnable, while the rest of the model’s parameters are frozen.
For the Llama2-7B model, the hyperparameters from Section 3.4 are used: $\tau=0.5$, $\gamma=2$, $\alpha=5$, and $\beta=0.4$. The model was trained for one epoch.

Metrics and Baselines.
The metrics considered include accuracy and the mean eviction ratio for all KV-Cache.
The accuracy of vanilla Llama2-7B on these datasets is used as the upper bound, with the goal of minimizing accuracy loss while maximizing the mean eviction ratio.
For KV-Cache eviction strategies, we select a local method (retaining only recent tokens) and StreamingLLM (Xiao et al., ) as representatives of static strategies, while H2O (Zhang et al., ) serves as a representative of methods based accumulative attention scores.
For all these methods, eviction ratio is set to 50%.

#### 4.1.2 Results

The results in Table 1 demonstrate the effectiveness of our method in balancing accuracy and KV-Cache eviction.
Our method consistently outperforms the baseline strategies in terms of accuracy across most tasks while achieving a higher mean eviction ratio.
This indicates that our method is able to evict more tokens on average without significantly compromising accuracy.

In terms of accuracy, our method achieves competitive performance, closely matching the performance of vanilla Llama2-7B in most cases.
For instance, on PIQA, COPA, and BoolQ, our method performs comparably to Llama2-7B, demonstrating minimal accuracy degradation.
Compared to baselines, our method shows superior performance, particularly in tasks like BoolQ and PIQA, where both accuracy and eviction ratio surpass the baselines. This confirms the advantage of learnable mechanism for more efficient token retention strategies without sacrificing model performance.

Moreover, it is worth highlighting that our method achieves these results with minimal computational overhead, as mentioned in Section 4.1.1.
The continual pre-training was conducted on only 5,000 samples and trained for just one epoch, demonstrating the lightweight nature of our approach.
This efficiency can be attributed to the fact that our method does not need to learn new knowledge from scratch but rather focuses on learning effective token retention strategies, leveraging the existing capabilities of the pre-trained model.

### 4.2 Supervised Fine-tuning

Figure: Figure 3: This visualization highlights attention patterns in Llama2-7B after fine-tuning on the BoolQ dataset, using a selected sample to showcase multiple heads within both MHA and AG of different layers. In part (i), we visualize attention scores from several MHA heads across different layers before eviction. 1. MHA heads exhibit diverse attention patterns, especially in the first two layers, where attention is highly heterogeneous across heads. 2. While the first two layers show dense attention patterns, the subsequent layers become progressively sparser. 3. Bright-yellow vertical lines consistently appear at the same position across nearly all heads, especially beyond the first two layers, indicating critical tokens for inference. This aligns with the *Heavy Hitters* in H2O (Zhang et al., ), where a small portion of tokens significantly contributes to the attention scores. Our method’s effectiveness is demonstrated by the persistence of these bright lines in deeper layers, showing that crucial tokens are retained rather than evicted, ensuring their importance is preserved across the network. In part (ii), we visualize the attention scores from the attention-like structure within the AG mechanism. As the model layers deepen, the attention pattern shifts from high-resolution to lower-resolution, indicating that AG increasingly focuses on distilling in-context information. Thus, deeper AG layers no longer need such high resolution to capture global information, as the model has already refined and summarized it. This suggests potential optimizations, such as reducing the number of heads or dimensionality in deeper AG layers, to further improve efficiency.
Refer to caption: x3.png

#### 4.2.1 Setup

Model and tasks.
We utilize Llama2-7B (Touvron et al., ) as base model.
To evaluate our approach, we select six widely recognized downstream tasks: PIQA (Bisk et al., ), ARC-C (Clark et al., ), RTE (Bar-Haim et al., ), COPA (Roemmele et al., ), BoolQ (Clark et al., ), and OpenBookQA (Mihaylov et al., ).
For each task, we fine-tune the model using the respective training set and evaluate its performance on the corresponding test set.

Baselines and Implementation Details.
For the baseline, we apply LoRA (Hu et al., ) to fine-tune the models on each task, targeting the $(W_{q},W_{k},W_{v},W_{o})$ weights in the self-attention modules.
In our method, in addition to applying LoRA to these parameters, the AG modules are also set to be learnable.
As mentioned in Section 3.4, the following hyperparameters are used: $\tau=0.3$, $\gamma=0$, $\alpha=1$ or $0.5$, and $\beta=0.28$.
We employ the AdamW optimizer (Loshchilov & Hutter, ) with a learning rate of 5e-5, and train for two epochs on each dataset.

#### 4.2.2 Results

As shown in Table 2, our method demonstrates a strong balance between accuracy and KV-Cache eviction across the six downstream tasks.

With $\alpha=1$, our method maintains competitive accuracy compared to the fine-tuned Llama2-7B baseline, with minimal drops (e.g., 82.15% vs. 82.92% on PIQA), while achieving a high mean eviction ratio of 60.00%.
With $\alpha=0.5$, the eviction rate is reduced to 55.49%. But in some tasks like RTE and COPA, our method even surpasses the baseline accuracy. The average accuracy also exceeds the baseline (79.06% vs. 77.86%), suggesting that effectively evicting redundant tokens allows the model to focus on relevant information and improves performance.

Additionally, under the same hyperparameter settings, the performance varies across tasks. For instance, ARC-C is more challenging to evict compared to OpenBookQA, leading to a larger accuracy drop post-eviction. This highlights the importance of task-specific KV-Cache eviction policies.

### 4.3 Visualization

**Table 3: Ablation study on various settings of AG, reporting accuracy (Acc.) and KV-Cache eviction ratio (%Eviction) under different configurations, with $\alpha=1$ in all settings. The results for (1) correspond to the setup described in Section 4.2, where the number of AG heads is 4, and the head dimension is 128. In (2-1) and (2-2), we explore the impact of the number of AG heads, with (2-1) using 2 heads and (2-2) using 1 head. The comparison between (1), (2-1), and (2-2) shows that reducing the number of heads leads to a drop in both accuracy and eviction ratio, indicating that the capacity of AG is closely tied to the number of heads. For (3-1) and (3-2), we assess the effect of reducing head dimensions for AG heads, where (3-1) has half the dimension size of (1) and (3-2) has 1/4. Comparing (1), (3-1), and (3-2) reveals that smaller dimensions reduce the eviction capability and accuracy, highlighting the importance of maintaining sufficient head dimensionality. Settings (4-1) and (4-2) examine using the hidden states and the AG module from the previous layer to inform the current layer’s eviction strategy. In (4-1), the first layer does not evict, and from the second layer onward, the eviction is determined by the previous layer. In (4-2), the first and second layers do not evict, following the observation from Section 4.3 that the first two layers are denser than subsequent layers. Both (4-1) and (4-2) show a slight decline in accuracy and eviction ratio compared to (1), but this approach introduces parallelism, offering a potential avenue for future optimizations. The setting (5) replaces the attention-like structure in AG with a simple linear layer to determine the eviction strategy. The comparison with (1) shows that linear layers almost cannot evict tokens effectively, reinforcing the necessity of leveraging global in-context information for successful eviction, as discussed in Section 3.3.**
|  | Metric | PIQA | ARC-C | RTE | COPA | BoolQ | OpenBookQA | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| (1) | Acc. | 82.15 | 59.66 | 64.26 | 93.00 | 86.82 | 78.80 | 77.45 |
| %Eviction | 66.16 | 48.31 | 65.47 | 45.40 | 67.46 | 67.17 | 60.00 |  |
| (2-1) | Acc. | 81.88 | 57.63 | 65.70 | 91.00 | 87.52 | 77.40 | 76.86 |
| %Eviction | 63.92 | 36.38 | 62.73 | 24.38 | 65.22 | 63.57 | 52.70 |  |
| (2-2) | Acc. | 82.15 | 53.90 | 62.45 | 89.00 | 87.31 | 77.40 | 75.37 |
| %Eviction | 58.97 | 31.47 | 59.77 | 20.32 | 63.02 | 59.17 | 48.79 |  |
| (3-1) | Acc. | 81.45 | 53.36 | 58.84 | 88.00 | 86.73 | 78.40 | 74.46 |
| %Eviction | 61.75 | 33.55 | 61.34 | 19.24 | 64.59 | 59.59 | 50.01 |  |
| (3-2) | Acc. | 83.03 | 53.90 | 59.93 | 89.00 | 87.16 | 76.40 | 74.90 |
| %Eviction | 58.68 | 24.23 | 32.23 | 12.28 | 59.40 | 55.54 | 40.39 |  |
| (4-1) | Acc. | 81.66 | 55.25 | 66.06 | 88.00 | 86.85 | 78.00 | 75.97 |
| %Eviction | 49.52 | 36.92 | 46.85 | 28.74 | 56.02 | 60.32 | 46.40 |  |
| (4-2) | Acc. | 82.75 | 55.93 | 79.06 | 82.00 | 86.33 | 78.40 | 77.41 |
| %Eviction | 53.31 | 44.38 | 51.20 | 47.95 | 61.98 | 61.73 | 53.43 |  |
| (5) | Acc. | 82.54 | 54.58 | 57.40 | 81.00 | 87.71 | 74.80 | 73.01 |
| %Eviction | 1.06 | 0.46 | 0.81 | 0.26 | 1.38 | 1.16 | 0.86 |  |

In this section, we present visualizations to highlight key characteristics of our AG mechanism.
After fine-tuning on specific tasks, we visualize the model’s MHA and AG weights for selected samples, as shown in Figure 3.
Additional visualizations can be found in Section A.2.

### 4.4 Ablation

In this ablation study, we explore the effects of various configurations on the AG mechanism, focusing on the number of AG heads, AG head dimensions, and eviction strategies.
For details of each setting and the corresponding performance, refer to Table 3.
Key findings include the impact of reducing the number of heads, as shown in (2-1) and (2-2), and the reduction of head dimensions in (3-1) and (3-2), both of which result in lower accuracy and eviction ratios.
In (4-1) and (4-2), we examine eviction strategies where the current layer’s eviction is based on the previous layer’s hidden states and AG module, which introduces parallelism but slightly reduces performance.
Finally, replacing the AG with a linear layer in (5) demonstrates the necessity of using an attention-like structure for effective token eviction.
For more ablation, please refer to Table 5.

## 5 Conclusion

In conclusion, the proposed Attention-Gate mechanism offers a flexible and adaptive solution to KV-Cache eviction in large language models.
By dynamically identifying and discarding less important tokens, Attention-Gate addresses the limitations of static and attention-score-based strategies, providing efficient context-aware eviction.
This mechanism integrates seamlessly with pre-trained models and can be easily tuned, making it a practical and effective method for enhancing both performance and memory efficiency in various tasks.

## Appendix A Appendix

### A.1 Results of Continual Pre-training on Mistral

We conducted continual pre-training on Mistral using 5,000 samples from RedPajama (Computer, ), and the results are shown in Table 4. Compared to the performance of Llama2-7B presented in Table 1, Mistral’s performance slightly declined. We hypothesize that this may be due to the distribution of RedPajama’s data being less suited to Mistral. Additionally, this raises the question of whether KV-Cache eviction is model-dependent, and whether its effectiveness is related to the model’s expressive power.
Although the parameter counts of Mistral-7B and Llama2-7B are similar, Mistral-7B significantly outperforms Llama2-7B. This could suggest that Mistral is utilizing more tokens or scoring them with finer granularity, which results in fewer redundant tokens and thus makes eviction less effective.

**Table 4: Performance comparison between Mistral-7B and Ours across various tasks.**
|  | Metric | PIQA | ARC-C | ARC-E | RTE | COPA | BoolQ | HellaSwag | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Mistral-7B | Acc. | 80.09 | 42.37 | 63.14 | 48.01 | 76 | 64.22 | 73.02 | 63.84 |
| Ours | Acc. | 75.9 | 34.24 | 55.2 | 48.01 | 65 | 62.2 | 67.91 | 58.35 |
| Eviction | 37.14 | 39.48 | 37.80 | 40.93 | 45.27 | 44.68 | 50.92 | 42.32 |  |

### A.2 More Visualization

Figure: Figure 4: The complete version of Figure 3.
Refer to caption: x4.png

Figure: Figure 5: More samples for the same scenario of Figure 3.
Refer to caption: x5.png

Figure 4 provides a comprehensive view of the layers and attention heads from Figure 3. Additionally, four sample scenarios for the same setup are presented in Figure 5.

### A.3 More Ablation

**Table 5: Exploring the impact of the number of recent tokens, viewed from the perspective of the attention matrix and considering slanted retention patterns. (1) corresponds to the setup described in Section 4.2, where only the current token is retained, and thus reflecting only the diagonal retention in the attention matrix. For (6-1) and (6-2), the number of recent tokens retained is set to 5 and 10, respectively. The results suggest that increasing the number of recent tokens does not necessarily enhance performance under the AG framework. Further exploration of how to manage recent tokens, such as applying learnable weighted strategies, could be an interesting direction for future work.**
|  | Metric | PIQA | ARC-C | RTE | COPA | BoolQ | OpenBookQA | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| (1) | Acc. | 82.15 | 59.66 | 64.26 | 93.00 | 86.82 | 78.80 | 77.45 |
| %Eviction | 66.16 | 48.31 | 65.47 | 45.40 | 67.46 | 67.17 | 60.00 |  |
| (6-1) | Acc. | 83.08 | 50.85 | 65.34 | 82.00 | 87.31 | 73.20 | 73.63 |
| %Eviction | 65.15 | 40.96 | 64.29 | 21.37 | 67.49 | 63.69 | 53.83 |  |
| (6-2) | Acc. | 81.61 | 53.56 | 60.29 | 82.00 | 87.37 | 74.20 | 73.17 |
| %Eviction | 65.66 | 44.48 | 65.14 | 24.28 | 68.18 | 63.44 | 55.20 |  |

In this section, we present additional ablation results in Table 5.

## Peer Reviews

confidence: 5
contribution: 1
initial_rating: 3
presentation: 3
questions: None
soundness: 2
strengths: - The design choices of the proposed method are clearly motivated.
- Presentation is clear and easy-to-follow
summary: This paper aims to reduce the memory overhead of the KV cache in LLM inference. The high-level goal is to ``evict’ ’certain tokens from the attention computation without hurting the model’s accuracy, since their key and value vectors (KV-cache) no longer need to be stored in GPU memory. Differently from previous works that find the evicted tokens through heuristics or attention pattern analysis, this work learns a binary gating function for each attention head that predicts whether a token should be attended to in a context-aware manner. Light training of these gating functions is needed.

Experiments with some tasks suggest that the proposed method can evict a substantial amount of tokens without hurting the model’s accuracy.
weaknesses: - Questionable choices of the evaluation benchmarks. The paper motivates the importance of KV-cache reduction by stating that “KV-cache can become a bottleneck when dealing with large models and long-context queries.” However, to the best of my knowledge, none of the benchmarks in this paper is “long-context,” nor are the models “large.” Besides, I am not sure benchmarks such as RTE and COPA are the best candidates for properly evaluating these models’ performance, compared to more “up-to-date” choices such as MMLU or BBH.
- Unconventional experimental settings. In the SFT experiment, the paper finetunes the model for each task, which was a common setting for previous-generation models such as BERT. For models such as Llama-2, it is more established to finetune the models with a mixture of a variety of datasets and evaluate the models’ general-purpose capabilities. Therefore, I am not sure whether the conclusions from this experiment would be useful for future research.
- Flawed comparisons and missing baselines. Most of the claimed gains of the proposed method, which requires training, are over zero-shot baselines. A more fair comparison would be to compare to baselines that use the same amount of training as the proposed method. Besides, these two baselines are missing: https://arxiv.org/abs/2308.16137 and https://arxiv.org/abs/2310.01801
- Unjustified claims about efficiency. The paper claims that the proposed method “demonstrate both efficiency and adaptability.” In my opinion, efficiency should be quantified using metrics such as memory overhead, latency, throughput, etc. The paper does not include any of such evaluations in its experiments.
- Limited results that make it unclear whether the proposed method can generalize. With limited selection of the base LLM and evaluation benchmarks, it is unclear whether the proposed method can be useful for models with a different size, architecture (e.g., group-query attention as in Llama-3 and Mistral), and more challenging tasks.
- Potential errors in the training objective: In Eq 1, the purpose of including $\beta$ in Eq 1 is unclear to me since it does have any impact on the gradients. This might be a typo and the paper might have missed a square of the bracket term.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: - Experiments are conducted on datasets with short input. In practice, the memory issue caused by the KV-Cache is severe in long-text evaluation settings. I would like to see the results reported on Longbench or Infinitebench datasets. This will make the AG more practical and widely adopted.
- I'm confused about the definition of Equation (1): The $\beta$ is task/layer-wise or not? As the authors illustrated in Figure (1), the eviction patterns changed across tasks and layers.
soundness: 3
strengths: - Attention-Gate is technically sound and easy to implement. The experiment showed its performance in both continued pre-training and lightweight fine-tuning settings with clear visualization of Attention Maps.
- Ablation studies are thoroughly conducted on the impact of the number/dimension of AG heads and the structure choice of AG.
summary: This paper introduces a simple parameterized KV-Cache eviction mechanism called Attention-Gate, designed to enhance the efficiency of LLMs inference by selectively discarding KV-Cache states for less important tokens. This approach addresses the limitations of traditional static eviction strategies, like those in StreamingLLM, and the attention bias issues in dynamic methods such as H2O. Attention-Gate operates by generating eviction flags for tokens based on the entire context. Experiments showed the performance could be retained with even high eviction ratio compared to baselines.
weaknesses: - The baselines like Local, StreamingLLM, and H2O are weak. As the motivation of Attention-Gate mechanism is to address the attention bias issue, some follow-up works like NACL and A2SF with the same motivation should be included as the baselines to fully show the effectiveness of AG.
- The additional training phase in Attention-Gate may influence the performance of the original LLama2-7B model, so the performance of LLama2-7B with continued training without AG mechanism should be reported.

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 2
questions: For the method, the following should be addressed.
1. Currently, it lacks an explanation of the attention-like structure of Attention-Gate, including the network design and computation flow. 
2. The only computation equation listed in the paper (Equation 1) is unclear and appears incorrect. 
* The definition and computation of $\bar{AG}_l$ are not explained.
* The necessity of introducing an additional threshold $\beta$ is unclear; its role seems to overlap with the probability threshold $\tau$.
* The direct subtraction of $\beta$ seems ineffective for preventing over-aggressive token eviction; an activation or truncation function may be missing.

More experimental analysis can be added.
1. Inference efficiency: including Attention-Gate before each MHA layer adds computational overhead. To assess inference efficiency, a comparison of the computation burden (in FLOPs) between the proposed method and traditional approaches is essential.
2. Comparative evaluation: as the paper mentions other learnable eviction methods like NACL and A2SF, their performance should be included for a more comprehensive comparison of effectiveness.
3. Eviction ratio variability: only a 50% eviction ratio is reported; assessing model performance across varying eviction ratios would provide a more comprehensive analysis.
4. Long sequence analysis: given that KV-cache eviction is critical for long sequences, it would be beneficial to include evaluation tasks involving longer sequences.
soundness: 2
strengths: 1. Compared to traditional static or local dynamic eviction methods, Attention-Gate offers a flexible and adaptive approach to KV-Cache eviction in LLMs, enhancing cache management by adjusting to contextual needs.
2. Attention-Gate opts for the attention-like structure to collect contextual information, which can vary among attention heads/layers and be seamlessly plugged into pre-trained LLMs.
summary: The paper introduces "Attention-Gate", a parameterized mechanism for KV-Cache eviction in large language models (LLMs). Unlike traditional static or local dynamic eviction methods, Attention-Gate uses contextual inputs to decide which tokens to discard adaptively. Experimental results on continual pre-training and supervised fine-tuning scenarios suggest that the proposed method can achieve better performance than traditional approaches.
weaknesses: 1. Limited architectural detail: The paper lacks a thorough explanation of the detailed design of Attention-Gate, particularly regarding its network architecture and computation equations.
2. Insufficient experimental analysis: The experiments would benefit from more comprehensive analysis to better illustrate the efficiency and effectiveness of Attention-Gate, providing a clearer understanding of its practical impact and performance in different scenarios.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: N/A
soundness: 2
strengths: - As far as I know, this paper is the first one that introduces parameterized module to achieve KV cache eviction. Previous studies usually rely on parameter-free heuristics.
- This paper is probably one of the first that achieves dynamic KV cache eviction at both layer-level and head-level. Previous studies belong to either layer-level dynamic KV cache eviction (e.g., PyramidKV) or head-level one (e.g., AdaKV).
- The design choices of Attention-Gate are thoroughly elaborated and sufficiently motivated.
- Experimental results show that Attention-Gate outperforms typical baselines with a similar KV cache budget.
summary: This paper is concerned about KV cache compression, especially KV cache eviction. This paper claims that previous parameter-free eviction strategies either can not adapt to the context or can only adapt to the biased context. To remedy this, this paper proposes a parameterized eviction strategy, which learns to evict KV cache at prefilling stage with a trainable module termed Attention-Gate. Experimental results on both continued pretraining and supervised finetuning demonstrate that, with similar eviction proportions, Attention-Gate can achieve better performance in comparison to existing eviction baselines such as StreamingLLM and H2O. Particularly, Attention-Gate combined with LoRA can even realize better performance than LoRA itself. Further ablations indicate that the design choices of Attention-Gate is adequate.
weaknesses: - This paper claims previous studies that can adapt to the given context (e.g., H2O) are largely biased and uses this as a core motivation of a parameterized KV cache eviction strategy (i.e., Attention-Gate). However, there is no clear evidence showcasing that Attention-Gate is not biased. Particularly, from Figure 1, we can see that there are still attention sinks in quite a few heads.
- The proposed strategy in this paper is limited to prefilling stage. There are indeed studies who are also limited to prefilling stage. However, it would be better to consider decoding stage especially in cases where the input prompt is comparably short and the output is long.
- The use of straight-through estimator is not very clear. I recommend more details on this.
- Though it is very nice to see that Attention-Gate combined with LoRA outperforms LoRA, it is not very clear why this would happen. And it is actually very interesting to uncover the reason why eviction can even lead to better results. Perhaps eviction can also serve as a sort of regularization?
- Frustratingly, while KV cache eviction is targeted at improving memory efficiency and potentially latency efficiency, this paper lacks an in-depth comparison on efficiency. And this would predominately affect the contribution and impact of this paper.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: Please refer to the Weakness.
soundness: 3
strengths: 1. The authors' method is structurally simple yet highly effective, offering more flexibility and better performance compared to previous eviction strategies. It incorporates global information and demonstrates significant improvements in performance.
2. This approach avoids the redundant attention computations seen in prior methods, enabling a more thorough and efficient eviction of tokens.
3. This method requires only a few thousand samples for post-training or fine-tuning to be effective, making it easily applicable to a wide range of existing pre-trained models.
4. The authors provide a detailed analysis of attention behavior across different layers, explaining the effectiveness of their approach and illustrating its specific impact across various cases.
summary: This paper introduces a dynamic KV-Cache eviction mechanism called Attention-Gate (AG) for large language models. AG selectively evicts less important tokens based on attention scores, unlike static strategies. Positioned before attention layers, AG can be seamlessly integrated into pre-trained LLMs. The method is validated on various tasks, demonstrating improved efficiency compared to static eviction strategies like StreamingLLM and H2O, with benefits for long-context queries.
weaknesses: 1. There is no detailed analysis of the data sources used for post-training if other pre-training data besides RedPajama is chosen. There is also no examination of how the amount of post-training data impacts the final results.
2. The method is only effective on Llama-2 and does not show improvement on Mistral-7B, but the paper lacks a detailed analysis of why this is the case. Additionally, there are no experiments on models of other scales besides 7B.
3. The paper lacks analysis on the inference efficiency improvement after eviction, as well as the additional time overhead introduced by the attention-gate computation compared to the original Transformer.
4. It would be helpful to detail the similarities and differences between MoD (https://arxiv.org/abs/2404.02258), which also uses a gate mechanism based on Softmax to select tokens for computation.
5. Most of the evaluation tasks are short-form, few-shot tasks, and there is a lack of comprehensive evaluation tasks that test various LLM capabilities (e.g., MMLU). Additionally, the inclusion of long-form datasets could provide more robust results.
6. There are several hyperparameters involved in AG, but there is no analysis of how these hyperparameters were selected or their impact on performance.
7. Appendices A.2 and A.3 are missing content.

Meta-review: This paper seeks to address the issue of kv-cache sizes growing fast with the size of models and input context. They do so by learning a model component that tries to determine which tokens in the past kv states need to be stored.

While the approach is simple and the eviction detection operates in a data-dependent way, the extent to which it works is a bit unclear. As highlighted by pretty much all reviewers, the choices of baselines as well as of evaluation benchmarks may be less than ideal, but besides that, I would highlight that the evaluation doesn't test for the main claim of the paper which is to enable reductions in memory footprint during inference. Eviction rates are reported, but the extent to which that actually yields savings in inference cost in space remains unclear. There are also some clarity issues, and the approach is itself not clearly defined as pointed out by some of the reviewers. I would encourage the authors to expand the evaluation section and refactor the intro sections with some more detailed description of the proposal and its justification.

Decision: Reject