NACL: A General and Effective KV Cache Eviction Framework for LLMs at Inference Time (2024)

Abstract: Large Language Models (LLMs) have ignited an innovative surge of AI applications, marking a new era of exciting possibilities equipped with extended context windows. However, hosting these models is cost-prohibitive mainly due to the extensive memory consumption of KV Cache involving long-context modeling. Despite several works proposing to evict unnecessary tokens from the KV Cache, most of them rely on the biased local statistics of accumulated attention scores and report performance using unconvincing metric like perplexity on inadequate short-text evaluation. In this paper, we propose NACL, a general framework for long-context KV cache eviction that achieves more optimal and efficient eviction in a single operation during the encoding phase. Due to NACL’s efficiency, we combine more accurate attention score statistics in PROXY-TOKENS EVICTION with the diversified random eviction strategy of RANDOM EVICTION, aiming to alleviate the issue of attention bias and enhance the robustness in maintaining pivotal tokens for long-context modeling tasks. Notably, our method significantly improves the performance on shortand long-text tasks by 80% and 76% respectively, reducing KV Cache by up to 5× with over 95% performance maintenance. The code is available at https: //github.com/PaddlePaddle/Research/ tree/master/NLP/ACL2024-NACL.

heading: 1 Introduction
text: Large Language Models (LLMs) with longer context window (Touvron et al., 2023; Xiong et al., 2023; Jiang et al., 2023; Anthropic, 2023; OpenAI, 2023) have emerged recently for better conducting long conversations, summarizing long documents,
∗Equal contribution. Work done at Baidu Inc. †Corresponding author. ‡ Project lead.
or debugging code at the repository level (Bai et al., 2023). However, their deployment is costly and infeasible on fixed memory hardware, mainly due to the surprisingly large memory consumption of KV Cache mechanism. For instance, a 7 billionparameter model with an input batch size of 4 and a sequence length of 32k results in 64GB of KV cache, 4.7× larger than the model weights.
To mitigate the pressure on the scarce GPU memory from using KV cache, a number of studies (Zhang et al., 2023; Liu et al., 2023c; Ge et al., 2023; Xiao et al., 2023) have explored sparsity among Transformer attention blocks to evict unnecessary tokens from the KV cache. For instance, H2O (Zhang et al., 2023) utilized the local statistics of accumulated attention scores to retain a balance of recent and heavy hitter tokens during generation. Window Attention based methods (Xiao et al., 2023) proposed to keep the initial tokens which is proven vital for generation fluency. This line of work reduced the memory footprint of KV cache
7913
for efficient inference with negligible loss in generation quality. In addition, the above methods do not require costly retraining which is more suitable for current open-sourced LLMs (Touvron et al., 2023; Jiang et al., 2023), compared to those that need specific attention mechanism adaptation (Beltagy et al., 2020; Kitaev et al., 2020; Shazeer, 2019; Ainslie et al., 2023).
However, we argue that the performance reported in the above methods is over-optimistic, as the evaluation metric and tested benchmark is not sufficient. LLMs may fail in real-life long-context modeling tasks (Bai et al., 2023), though they can achieve low language modeling perplexity which is untrustworthy used as the golden metric in current studies (Xiao et al., 2023; Han et al., 2023). Furthermore, the local statistics of accumulated attention score is observed to be biased (see Fig. 2), especially in long context input, meaning that it should be carefully used as the only strategy for measuring the importance of tokens.
To fill the gap, we propose NACL, a general and effective KV cache eviction framework to unleash the power of LLMs for long-context modeling with limited memory budgets. NACL specifically formulates the eviction task in encoding phase which is different from the commonly used one-tokenin one-token out eviction procedure in generation phase. In encoding phase, the eviction can be effectively implemented to apply only once on the whole input by progressively evicting KV caches layer by layer. The one-eviction formulation benefits current eviction policies in a more efficient and optimal way, as multiple costly eviction operations can be combined, then the global statistics of attention scores can be utilized.
Based on the above formulation, we present PROXY-TOKENS EVICTION which exploits the global statistics of attention scores gathered from proxy tokens for eviction. In practise, the proxy tokens can be selected from the question input, commonly located at the end of a long text. Intuitively, these proxy tokens are more capable of retaining the task-specific tokens in KV cache. As a result, PROXY-TOKENS EVICTION alleviates the attention bias problem (see Sec. 4) occurred in methods using local statistics (Zhang et al., 2023; Oren et al., 2024) or task-irrelevant proxy tokens (Liu et al., 2023c).
However, PROXY-TOKENS EVICTION also relies heavily on the statistic of attention scores
which may be untrustworthy in long-context input. Thus, we incorporate RANDOM EVICTION, a random eviction policy, into PROXY-TOKENS EVICTION. RANDOM EVICTION randomly samples tokens to evict from the probability distribution in PROXY-TOKENS EVICTION with different seed on attention heads and layers. This diversified randomness enhances the model’s robustness to maintain potentially important tokens in long text generation.
We conducted extensive experiments on a single NVIDIA A100 (80GB) GPU on representative open-sourced LLMs: LLaMA2-base, LLaMA2Chat (Touvron et al., 2023), and evaluated them on both short- and long-text modeling tasks from lmeval-harness (Gao et al., 2021) and LongBench (Bai et al., 2023). The experiments show that NACL performs KV cache eviction efficiently with negligible degradation on model quality (i.e., saving the inference memory usage of KV cache by up to 5× with over 95% maintenance). Specifically, NACL achieve 80% and 75% performance improvement on short- and long- text modeling tasks, respectively, with 50% KV cache reduction, compared to current eviction methods.

heading: 2 Related Work
text: Efficient Inference with Limited KV Cache Budgets emerged for reducing the prominent inference bottleneck caused by KV cache, particularly for long content input. A series of methods (Zhang et al., 2023; Ge et al., 2023; Liu et al., 2023c; Oren et al., 2024) explored the sparsity among Transformer’s attention block, then evicted unnecessary tokens from KV Cache for efficient inference. For instance, H2O (Zhang et al., 2023) retained a balance of recent and heavy hitter tokens with the highest accumulated attention scores throughout the sequence. Scissorhands (Liu et al., 2023c) sequentially predicted the potentially pivotal tokens with the attention score above average within a history window. Some method (Ge et al., 2023) further applied costly eviction policy selection for better performance. However, the above methods relied heavily on the attention score with local statistics which may be sub-optimal in long-context tasks (Dacheng Li*, 2023; Bai et al., 2023).
Meanwhile, some efforts have been made to utilize a learnable mechanism to determine necessary tokens during inference (Anagnostidis et al., 2023), or converting the traditional multi-head at-
tention(MHA) (Vaswani et al., 2017) to multi-query attention (MQA) (Shazeer, 2019) or group-query attention (GQA) (Ainslie et al., 2023). However, these methods involve additional training, while NACL focuses on the inference phase without resource-intensive training.
Efficient Transformers (Tay et al., 2020) have been extensively explored (Child et al., 2019; Kitaev et al., 2020; Zaheer et al., 2020; Beltagy et al., 2020; Dai et al., 2019; Ding et al., 2020; Bulatov et al., 2022; Chevalier et al., 2023) to address the self-attention operation which scales quadratically with the sequence length. For instance, Sparse Transformer (Child et al., 2019) uses a dilated sliding window the reduces the attention complexity. Longformer (Beltagy et al., 2020) and Bigbird (Zaheer et al., 2020) reduced the complexity of selfattention by combining random, window and global attention. Recurrence Transformers (Dai et al., 2019) maintain a memory bank of past KV cache to process the long text in segments. However, the above methods either trade off model quality or require re-training of models, but often failed in achieving memory saving and wall-clock speedup at inference time (Dao et al., 2022).
Length Extrapolation enabled language models to generalize beyond the context window they were trained on. A recent line of research (Chen et al., 2023; Peng et al., 2023; Liu et al., 2023b) focuses on adapting relative positional embedding (Su et al., 2024) widely used in current Foundation models (Touvron et al., 2023; Jiang et al., 2023) for context window extension. Attention Sink (Xiao et al., 2023) and LM-Infinite (Han et al., 2023) further exploited the initial tokens to recover the performance of window attention for infinite-length inputs. However, the ability of these methods tested using metric like perplexity is over-optimistic for long context tasks (Dacheng Li*, 2023; Bai et al., 2023).

heading: 3 Problem Formulation
text: This section defines a two-phased approach for efficient KV cache management during LLM inference, tailored for scenarios with limited KV cache budgets.
Eviction Policy We defined the eviction policy Fscore : S i t ← Sit−1, subject to |Sit | = |Sit−1| ≤ C where the scoring function Fscore assigns low scores
to unnecessary tokens for eviction, such that the pre-define KV cache budget C is maintained. Sit denote the indices set of retained tokens in KV cache at t-th time step and i-th transformer layer.
Encoding Phase Eviction The model processes the input prompts, xiprompt = [x i 1, . . . , x i p] ∈ Rp×d , to compute the initial key cache Ki0 = xipromptW i K ∈ Rp×d and value cache V i0 = xipromptW i V ∈ Rp×d, where p denotes the encoding prompt length, W iK ,W i V ∈ Rd×d represent the key and value projection weight at layer i with hidden dimension d, respectively. The attention scores Aiprompt ∈ Rp×p are computed as (xipromptW i Q)·(xipromptW iK) T
√ d where W iQ ∈ Rd×d represent the query projection weight. The eviction in encoding phase is defined as follows:
Siencoding = Fscore(A i prompt, C)
then, the initial KV cache can be updatedKi0,V i0 ← KSiencoding ,VSiencoding for the later usage in generation phase.
Generation Phase Eviction Denote the generated tokens’ input to i-th layer as xidecoding = [zi1, . . . , z i T ] ∈ RT×d. The Generation phase updates the KV cache with each new token generation. Given the time step t and layer i, key and value cache is updated as Kit = [Kit−1, zit ·W iK ], V it = [V it−1, zit ·W iV ], respectively. The attention scores Ait ∈ R1×|K i t| are computed as (zitW i Q)·Kit T
√ d . The eviction in generation phase is defined as follows:
Sit = Fscore(A i t, St−1, C)
where the KV cache are updated Tm times following Kit,V it ← KSit ,VSit at every m time steps.
To note that, recent works formulated the encoding phase eviction the same as the one in generation phase which require step-by-step evictions, resulting in computational overhead. In contrast, we formulate the eviction to perform only once during the encoding phase and Tm times during the generation phase. Generally, the condition T ≪ p is readily satisfied in long text scenarios, allowing Tm to be approximated as a constant order of magnitude. Consequently, the overall time complexity is reduced from O(p+ T ) to O(1). This also allows the eviction policy in a global optimal manner comparing to those greedy algorithm that couples the input window size with the KV cache budget.

heading: 4 Observation
text: We present two experimental findings by rethinking previous eviction methods that inspire the design of NACL.
Rethinking Evaluation Metrics for Long-text Eviction Strategy Current metrics such as the perplexity (PPL) fall short in capturing the nuances of model performance in long-text scenarios, revealing a gap between evaluation practices and real-world applications (see, Tab. 1). Evaluations predominantly utilize datasets with short texts, inadequately representing the complexities and challenges of processing and understanding long-text input. The emphasis on textual fluency leads to a notable bias: the method (Xiao et al., 2023), though claiming for infinite input, fails in tasks (see, Tab. 1, 2) which requires the ability to generate accurately. This inspires us to re-evaluate current methods on both short- and long- text modeling tasks demands on comprehension and generation capabilities.
Rethinking Attention Scores to Retain Pivotal Tokens Attention bias problem refers to the phenomenon where, at each step of generation, attention scores are higher within the tokens directly preceding the current token, while comparatively diminished for all others. In Fig. 2 (a) and (b), the attention bias problem is observed, leading to an overemphasis on either initial tokens (Zhang et al., 2023) or recent tokens (Oren et al., 2024), overlooking those potentially pivotal tokens in longer context. Furthermore, the attention score distribution become flattened with the increase in text length (see Fig. 2 (d)), which may be less capable of accurately identifying important tokens. Normalization can solve this problem to some extent, but as stated in the H2O (Zhang et al., 2023) , the effect is not optimal. This inspires us to reform
the attention-based eviction methods to be less bias and more robust in long-context modeling tasks.

heading: 5 NACL
text: In this section, we present a hybrid KV cache eviction policy in NACL, including the PROXYTOKENS EVICTION in Sec. 5.1 and RANDOM EVICTION in Sec. 5.2.

heading: 5.1 Eviction based on Proxy Tokens
text: Based on previous observations, the current Fscore of accumulating attention scores effectively identifies important tokens but suffers from significant bias. We attribute this to the excessive redundant information in the process of scoring tokens.
We discovered that when calculating the attention for a given token x (i.e. tokens may need to be evicted), only a mere fraction of tokens xp (i.e. proxy tokens) are responsible for yielding the most precise outcomes during the computation of the token score. Hence, we introduce the proxy tokens hypothesis: within the input xprompt, there exists a subset called proxy tokens P ∈ xprompt, which precisely estimate the importance of tokens. Scoring function Fscore can be instantiated as:
Fscore (A, C) = ∑
xp∈P Softmax (A(xp, ∗))
The Fscore , calculated by reducing the attention score matrix column-wise using the proxy tokens subset P , provides the most precise measurement of a token’s importance during the eviction process. We can validate the significance of proxy tokens using a straightforward approach. When the proxy token is set to the universal set, our method is equivalent to H2O (Zhang et al., 2023), introducing redundant information that degrades the quality of eviction. When the proxy token is set to only the current token, our method can be equivalent to
MSRNN (Oren et al., 2024), neglecting substantial information, thus reducing the accuracy of eviction.
Due to the progressively flattened distribution of attention scores with the increase in text length, the pre-defined threshold for sampling Cp results in lack of generalizability for long text tasks. Therefore, we model the KV cache eviction as an optimization problem, aiming to find a set St that maximizes the function Fscore, while satisfying the constraint |St| = Cp.
St ← (argmax St⊂R
∑ x∈St Fscore(A, Cp)) ∪ P
where R = xprompt\P as the proxy tokens are retained by default. In practise, the proxy tokens tend to be chosen at the end of the input where the user’s question with more task-specific information is located in. The choice of Proxy Token can be based on task orientation, which allows our approach to be flexibly adapted to various application scenarios. For more information, please refer to the Appx. A.3.

heading: 5.2 Eviction based on Random Possibility Sampling
text: Eviction algorithms commonly rely on the attention scores which may be biased or lack robustness in capturing critical information throughout the generation process. Herein, we introduce a simple and effective eviction policy which incorporates the randomness into the decision-making process of the attention mechanism. By randomly sampling from a probability distribution, our method aims to enhance the model’s ability to recover and maintain important information that might otherwise be lost.
In detail, we construct the probability distribution from Frandom . Frandom can signify each candidate token’s relative significance in long text generation, and the probability Pprompt is determined as follows:
Pprompt = Softmax ( Frandom (Aprompt , Cr) )
where Pprompt allows the non-deterministic selection of pivotal tokens. Through this probabilistic lens, our model casts the dice, diversifying its focus and increasing the chances of preserving essential information across the span of long texts. Thus, we present RANDOM EVICTION with budget Cr:
Srandom ∼ Pprompt, |Srandom| = Cr
In practise, Prandom can be based on the normalized distribution Softmax(Fscore), then the complexity is O(|xprompt|) dominated by the softmax operation.
Finally, NACL effectively combines PROXYTOKENS EVICTION and RANDOM EVICTION, applying an efficient one-eviction strategy under the KV Cache budget C = Cp + Cr, shown in the following Algorithm 1. Our method is compatible with FlashAttention-2 (see Appx. A.7) to minimize memory and computational overhead, helping models to be efficiently deployed in long text tasks.

heading: 6 Experiments

heading: 6.1 Setup
text: Objective We aim to provide experimental evidence for three key research questions: 1. Whether there are advantages in performance and task generalization of NACL over other eviction methods. 2. How the two eviction policys in NACL affect the
Algorithm 1 NACL Algorithm 1: Total Cache budget C (C = Cp + Cr), Proxy-Token Evic-
tion Cache budget Cp, Random Eviction Cache budget Cr, Proxy tokens P , KV Cache K,V
2: function ENCODING(Prompts) 3: for Every Layer-i in LLMs do 4: for Every Attention Head-n do 5: W i,nQ ,W i,n K ,W i,n V ∈ Rd×d 6: A← (xpromptW i,nQ ) · (xpromptW i,nK ) T√ d −1
7: Fscore = ∑
xp∈P Softmax (A(xp, ∗)) 8: R← xprompt\P i,n 9: uscore ← (max ∑ x∈R Fscore(A, Cp)) ∪ Pi,n
10: urandom ∼ Softmax (Frandom(Aprompt), Cr)) 11: Si,nencoding ← uscore ∪ urandom 12: end for 13: end for 14: end function 15: function GENERATION(Sencoding, Max Length) 16: m← eviction interval 17: z0 ← last prompt token 18: S0 ← Sencoding 19: for t ∈ {1, ...,Max Length} do 20: for Every Layer-i in LLMs do 21: for Every Attention Head-n do 22: Ki,nt−1 ← KSi,nt−1 ,V i,n t−1 ← VSi,nt−1 23: Ki,nt ← [Ki,nt−1, zit ·W i,nK ] 24: Vi,nt ← [Vi,nt−1, zit ·W i,nV ] 25: A = (ztW i,nQ ) · Ki,nt T√ d −1 26: if t mod m = 0 then 27: Si,nt ← Eviction(A, C) 28: ▷ Ref: Line7-10. 29: end if 30: end for 31: end for 32: zt ← sample from LLM prediction 33: end for 34: end function
final functionality, and by what combination can we achieve optimal results. 3. What is the rationale behind NACL for superior results?
Models and Tasks We use the family of decoderonly Transformers: LLaMA2-7B-base, LLaMA27B-Chat (Touvron et al., 2023) to evaluate the effectiveness of NACL. To evaluate the few-shot learning ability, we sample seven tasks from the popular benchmark (lm-eval-harness (Gao et al., 2021)): PiQA (Bisk et al., 2020), COPA (Roemmele et al., 2011), OpenBookQA (Mihaylov et al., 2018), Winogrande (Sakaguchi et al., 2021), SciQA (Welbl et al., 2017), ARC-E and ARCC (Clark et al., 2018). In the long text scenario, we took seven tasks from Longbench (Bai et al., 2023): PassageRetrieval-Zh, PassageRetrieval-En, RepoBench-P (Liu et al., 2023a), HotpotQA (Yang et al., 2018), NarrativeQA (Kočiskỳ et al., 2018), TriviaQA (Joshi et al., 2017), QMSum (Zhong et al., 2021). We report perplexity computed on the OpenBookQA dataset as a measure of the model’s generation ability in generalized domains. We conduct our experiments on a single NVIDIA A100 80GB GPU. Results were averaged over various seeds to ensure reliability.
Baselines We consider four representative eviction methods:
• Attention Sink (Xiao et al., 2023) keeps the initial and recent tokens for infinite-length text processing.
• H2O (Zhang et al., 2023) firstly proposes utilizing the summation of attention scores for greedy eviction, which achieves fair results which serves as our main baseline.
• MSRNN (Oren et al., 2024) considers the current token’s attention score for eviction.
• Scissorhands (Liu et al., 2023c) increments the counter within a history window for low score
token eviction.

heading: 6.2 Result
text: Short-Text Performance Our experimental analysis shows the effectiveness of NACL in managing KV cache under constrained memory budgets while maintaining high performance across various shorttext benchmarks. Firstly, NACL demonstrated superior performance in comparison to the baseline eviction methods with minimal performance degradation. In the 5-shot setting, NACL achieved an average score of 63.8% points, nearly matching the full cache performance of 64.6% points and significantly outperforming H2O by 3.5% points. Moreover, NACL exhibited consistent improvements across most datasets relative to previous methods, affirming its robust generalization and practical applicability. In the 25-shot setting, although there was a performance dip across all methods due to the increased complexity and information redundancy, NACL still showed remarkable resilience. Notably, it matched or slightly outperformed the full cache setup in certain datasets, such as maintaining a 87.0% point score on COPA, identical to the full cache performance. This illustrates that NACL not only manages to select pertinent information effectively but also mitigates the impact of redundant data, enhancing the model’s robustness.
Long-text Performance NACL achieves 80% memory usage reduction with only mere 0.7% point decrease with respect to the average accuracy in Tab. 2. Fig. 5 (Left) shows that NACL is possible to achieve 3× more reduction in KV Cache while maintaining comparable performance to baselines. Additionally, we observed the stable performance of NACL under different budget, while others’ fluctuate. In HotpotQA and QMSum, NACL (30%) even surpassed the performance without KV cache eviction by 0.2% and 0.9% points, re-
spectively. For challenging passkey retrieval tasks, H2O and MSRNN with the attention bias towards initial and recent tokens fails in retaining the pivotal passkey located in the middle of the long input. In contrast, NACL demonstrates stable and superior performance in different budgets setting, that only missed 2 passkeys in PR-Zh and PR-En comparing to the model in full cache setting. This remarkable achievement highlights NACL’s ability to retain essential information, avoiding the pitfalls of redundant data and thereby bolstering the model’s robustness in processing complex long texts.

heading: 6.3 Ablation Studies

heading: The Effect of PROXY-TOKENS EVICTION
text: Proxy tokens play an important role in finding pivotal tokens. The performance degradation (see Tab. 3) is significant when removing this policy. In Fig. 5 (Middle), we report the impact of proxy token budget on the average accuracy as a proportion of the text length. In extreme cases, such as 0% and 100% proxy token budget, the method degenerates into two special cases: MSRNN (Oren et al., 2024) and H2O (Zhang et al., 2023), respectively. The suboptimal performance with 0% proxy token budget suggests that the unsufficiency of a single current token for determining the pivotal tokens. However, excessive abuse of proxy token budget up to 100% will introduce redundant information leading to decline in performance. In practise, we suggest the budgets for proxy tokens ∼10% for better performance.
The Effect of RANDOM EVICTION As shown in Tab. 3, RANDOM EVICTION obtained a performance gain of 9% points. In Fig. 5 (Right), increasing the random budget from 10% to 70% results in an average performance improvement of 2.25% to 8.17% points over 0% random budget. The performance peaks at 70% budgets, demonstrating the necessity of the RANDOM EVICTION. However, when the random budget’s proportion increases from 90% to 100%, a noticeable performance decline occurs, highlighting the importance to combine the Attention-score-based, PROXY-TOKENS EVICTION.
The Choice of Sampling Distributions for RANDOM EVICTION In Tab. 3, the experimental results demonstrate that sampling based on global statistical attention scores outperforms those based on uniform distributions in terms of performance. This indicates that attention scores can also provide
a more informative reference for randomness in specific contexts.
The Effect of Global Eviction The one-eviction formulation in NACL enhanced the average performance of 1.4% points in Tab. 3. Compared to the greedy algorithm, our approach reduces complexity while giving more consideration to global information. Furthermore, our algorithm exhibits greater simplicity and directness in its engineering implementation.
The Effect of Head-wise Token Eviction The results in Tab. 3 show a gradual decline in the algorithm’s effectiveness as the strategies become more uniform. The algorithm performs best when each head adopts a completely different strategy. The diversity of strategies leads to improved gener-
alization, preserving information across a broader spectrum of dimensions.

heading: 6.4 Analysis
text: Memory Usage of KV Cache In Fig. 4, we report the memory reduction of KV cache for LLaMA2-7B with respect to the sequence length with a fixed batch size of 4 in bf16 precision. At the same accuracy as H2O (20%), our method alleviates the linearly growing KV cache to 10% of the original size, significantly reducing the memory footprint of the KV cache for all model sizes. The memory usage reduction is more significant on long texts, effectively alleviating the memory bottleneck problem of long text reasoning.

heading: Interpretable Analysis of Eviction Results
text: Fig. 2 shows the attention score matrix where the darker area represents the retained tokens after eviction. Compared to previous methods, the proxy token guides NACL to sample the middle tokens more evenly, and protects the initial and recent tokens in the meantime. The head-wise randomness enables maintaining more context information, thereby enhancing the robustness of NACL.
Why Head-wise Eviction Matters From a probabilistic perspective, it is basically impossible for a token to be evicted in the head-wise eviction setting. Taking LLaMA-7B with 32-layers (number of layers l) and 32-heads (number of heads h) as an example, the probability of a token retained in least one head’s KV cache is 1−Ch which equals 99.92% when the KV Cache budget C = 20%. Even in a severe KV Cache budget setting like C = 1%, the probability that the information of a token is retained in at least one layer is 1 − (Ch)l which is larger than 99.99%.

heading: 7 Conclusion
text: In this paper, we focus on the accuracy, robustness, and reliability of evaluation for KV cache eviction algorithms deployed in LLMs for processing long texts. We introduce NACL, a novel approach that combines PROXY-TOKENS EVICTION and RANDOM EVICTION for KV cache eviction strategies, significantly reducing memory usage during model inference without the need for training. We model the eviction problem as a combinatorial optimization issue, where PROXY-TOKENS EVICTION provides eviction references based on importance, and RANDOM EVICTION enhances information richness and robustness through headwise and layerwise composite sampling. Through extensive evaluation, we demonstrate that NACL can significantly improve cache eviction strategies, reduce inference memory costs, and minimize the impact on the LLM’s ability to handle complex tasks.

heading: Limitations
text: Our approach presents two main limitations: First, due to constraints on resources, our method has not been extensively tested across various large-scale language models, especially for different lengths and even ultra-long texts. However, based on our current comprehensive experimental conclusions, we believe NACL can be extended to more application scenarios. In addition, we introduced the utilization of proxy tokens in PROXY-TOKENS EVICTION for identifying pivotal tokens, yet the selection of proxy tokens primarily relies on observations and experience. Determining proxy tokens from the model adaptively and accurately presents a challenge, which we deem worthy of further research.

heading: Ethics Statement
text: In this research, we employ open-source data and technologies, significantly reducing privacy concerns. Our innovative approach is geared towards understanding model contexts and boosting inference efficiency, with the aim of developing accessible and highly efficient models for extended contexts. This strategy is anticipated to propel the openness of NLP technology and its practical implementation in diverse applications. Importantly, our method is designed to be independent of the training process, ensuring it does not perpetuate or introduce biases into models. By focusing on cutting-edge and resource-efficient methodologies, our work contributes to making AI more open and automated, pushing the envelope in artificial intelligence while ensuring the benefits of our advancements are widely accessible and applicable across various domains, marking a step towards a more inclusive and automated AI future.

heading: Acknowledgments
text: We thank the anonymous reviewers for their insightful comments and constructive suggestions. We would like to thank Yinqi Yang, Jiawei Sheng, Xinhua Zhang, Shicheng Wang, Chuanyu Tang and members of the IIE KDsec group for their valuable feedback and discussions. We thank Siming Wu for the implementation of Reduce Attention Scores CUDA kernel. Work done during Yilong Chen’s internship at Baidu Inc. This research is supported by the National Key Research and Development Program of China (Grant No.2021ZD0110501 and Grant No.2021YFB3100600) and the Youth Innovation Promotion Association of CAS (Grant No.2021153).

heading: A Appendix

heading: A.1 Sparsity in Attention Cache
text: Inspired by previous literature on the existence of attentional sparsity in self-attentive heads, we delve into the sparsity of attention during the generation of LLMs. Given a normalised attention score matrix calculated by the softmax function applied to the dot product of the query matrix Q and the key matrix K, represented as A = Softmax(QK⊤), the attention mechanism allocates weights to different elements in the input sequence, reflecting their relative importance. Thus if the attention score of a token is low, it means that it has little influence on the process, and therefore we base our sparsification on the threshold to quantify sparsity. The sparsity percentage for a given threshold t is calculated as:
Sparsity(t, i) =
∑N j=1 1 (|Aij | < t)
N
where N is the dimension of the attention matrix, Softmax(QK⊤)ij represents the attention weight between the ith and jth elements, and 1(·) is an indicator function that returns 1 if the condition is true and 0 otherwise. This formula calculates the proportion of attention weights that are considered negligible or insignificant for different sparsity thresholds, thus providing a multi-faceted view of attention distribution’s sparsity across the model.
Fig. 6 shows the change of Sparsity under different thresholds, it can be seen that under different threshold values, the sparsity of A gradually increases with the sequence length, when the sequence length is 1200, the sparsity is more stable at 0.78. This means that 22% of the tokens are the dominant factor in the computation process.

heading: A.2 Attention Score Function in Eviction
text: In this section we show the importance scores of token in the contexts for H2O (see Fig. 7), MSRNN (see Fig. 8) and NACL (see Fig. 9) during different steps. The importance score function of H2O assigns larger scores to tokens in the front part, and MSRNN assigns larger scores to tokens in close proximity.NACL evenly distributes importance scores over longer contexts, so that both distant and close tokens will have a chance to be sampled to be retained.

heading: A.3 The Influence of Proxy Token Locations on NACL’s Proxy Eviction Strategy
text: In Proxy Token Eviction, we need to find the most accurate tokens in the sequence that can calculate the importance score of the token.The core of this problem is how to establish a judgment criterion of whether a token is important or not, and this judgment criterion determines how to select the proxy token set. Intuitively, we believe that whether a token in a sequence is important or not is determined by the task the model is about to accomplish. We refer to this task of the model’s input as the user’s question. In practical applications, we can usually separate the user’s question, so that we can place it at the end of input to maximize the performance of the Proxy Eviction Strategy. When we are unaware of the position of the user’s question, we primarily utilize proxy tokens to protect the beginning and end of the sequence, as in most cases, these positions contain crucial information about the generation. Even if the proxy token fails to include any information related to the question, our method can be considered as an improved version of MSRNN (Oren et al., 2024). By introducing the proxy token, we regularize the distribution of the Important scores. We also enhance the robustness of preserving intermediate information by combining it with the Random Sample strategy. Based on the aforementioned combination strategies, NACL demonstrates significant performance improvements in terms of results.

heading: A.4 Why NACL ues Random Eviction?
text: We demonstrate a case of combinatorial optimization based on randomness in Figure 10. Within the KV cache, the loss of any information can lead to a misunderstanding of the current information.
In previous methods, a broader spectrum of information that might otherwise be overlooked due to the uniformity and potential bias of score-based eviction methods. This randomness of diversification ensures that even information that is not prominently featured by self-attention mechanisms is given a chance to be integrated at deeper levels of the model (Dai et al., 2019). Consequently, this approach facilitates the model’s ability to capture and process a more varied set of token interactions, enhancing its overall performance by reducing the risk of losing vital information.

heading: A.5 Algorithmic Complexity Analysis
text: In long-text scenarios, the length of the generated text during the generation phase is often much shorter than the length of the input text. Assuming the length of the input text during the encoding phase is p and the length of the output text during the decoding phase is T , where T ≪ p. Given the complexity of the Fscore function as O(p), and applying a step-by-step eviction algorithm on data of length p results in a complexity of O(p2), making it impractical and costly for real-world applications.
Therefore, in long-text scenarios, we employ a one-time optimal eviction algorithm, where we calculate the optimal eviction strategy Senocding in a one-time computation during the encoding phase. Since the number of tokens generated is negligible compared to the total, we apply the eviction strategy Si at this stage as well. In comparison to the greedy algorithm that evicts based on Si−1 while retaining the same budget, our method can globally optimize to find the best eviction strategy. Moreover, we can decrease the time complexity from O(p2) to O(p), making the algorithm straightforward and effective in engineering applications.

heading: A.6 Hyperparameters
text: Here, we provide the hyperparameters for allocating the ratio of KV cache bugets for the hybrid eviction policy used in our experiments in Tab. 4.

heading: A.7 Reduce Attention Scores with FlashAttention-2
text: We have implemented NACL on 128k long-text inference and made it compatible with Flash Attention. There are two implementations below, both of which can be directly used with Flash Attention.
Re-Computation of the Attention Score: We utilize QP and K to calculate the required atten-
tion scores during the encoding phase, separate from Flash Attention. Since the proxy tokens set is very small, only a small portion of the attention score needs to be re-computed, thus the additional overhead is insignificant. According to experimental results, on a 128k context, evicting 20% while maintaining a stable 15GB of memory usage does not affect the inference speed.
Implementation of Reduce Attention Scores Kernel: The forward computation of FlashAttention2 (Dao et al., 2022) returns the log-sum-exp (Logsumexp) for each row. Leveraging this Logsumexp, we can recompute the attention scores matrix in the manner described in the backward computation of FlashAttention-2. Subsequently, we perform a column-wise summation to obtain the reduced attention scores, as outlined in Algorithm 2.
Budget Protect Proxy No-protect Proxy Proxy-Tokens Eviction Random Eviction
10% 1% 3% 2% 7% 20% 2% 18% 6% 12% 30% 1.5% 20% 10.5% 18%

heading: Future Work
text: By focusing on cutting-edge and resource-efficient methodologies, our work contributes to making AI more open and automated, pushing the envelope in artificial intelligence while ensuring the benefits of our advancements are widely accessible and applicable across various domains, marking a step towards a more inclusive and automated AI future