LagKV: Lag-Relative Information of the KV Cache Tells Which Tokens Are Important (2504.04704v1)

Abstract: The increasing size of the Key-Value (KV) cache during the LLMs long-context inference is the main obstacle for its balance between the deployment cost and task accuracy. To reduce the KV cache size in such scenarios, most previous efforts leveraged on the attention weight to evict non-critical cache tokens. But there is a trade-off in those methods, they usually require major modifiation of the inference infrastructure and significant computation overhead. Base on the fact that the Large Lanuage models are autoregresssive models, we propose {\it LagKV}, a KV allocation strategy only relying on straight forward comparison among KV themself. It is a totally attention free method which offers easy integration to the main stream inference platform and comparable performance comparing to other complicated KV compression methods. Results on LongBench and PasskeyRetrieval show that, our approach achieves nearly zero loss when the ratio is $2\times$ and $\approx 90\%$ of the original model performance for $8\times$. Especially in the 64-digit passkey retrieval task, our mehod outperforms the attention weight based method $H_2O$ over $60\%$ with same compression ratios. Our code is available at \url{https://github.com/AI-Lab-China-Merchants-Bank/LagKV}.

An Overview of "LagKV: Lag-Relative Information of the KV Cache Tells Which Tokens Are Important"

The paper by Liang et al., titled "LagKV: Lag-Relative Information of the KV Cache Tells Which Tokens Are Important," addresses a prominent challenge in the deployment of LLMs related to the size of the Key-Value (KV) cache during long-context inference. As the complexity of tasks increases, so does the KV cache size, which subsequently raises the computational and deployment costs. Traditionally, methods to manage KV cache have relied on attention mechanisms, which, while effective, bring significant computational overhead and are often incompatible with certain optimization strategies like Flash Attention (FA). This paper presents LagKV, an innovative KV allocation strategy that works independently of attention weights and instead utilizes token and channel-wise distribution patterns within the KV space, offering a more practical and efficient approach.

Core Methodology

LagKV capitalizes on the autoregressive nature of LLMs by proposing a KV compression approach void of attention mechanisms. The method hinges on the "token-wise locality" property, which suggests that tokens in close proximity exhibit more similar Key/Value tensor values compared to tokens further apart. The compression approach operates by dynamically applying a compression after the prefill stage. It always maintains a smaller subset of the attention sink and partitions the remaining KV cache into lag-sized blocks, where KV scores are computed based on min-max normalization against subsequent partitions. This scoring process uses the channel-wise standard deviation of the normalized Key and Value states to determine token importance, which is then employed to prune tokens using a Top-k strategy.

Experimental Results

The experimental analysis utilizes benchmarks such as LongBench and the challenging Passkey Retrieval Task. Key results highlight that LagKV achieves nearly zero loss at 2x compression ratios and maintains ~90% of original model performance at 8x compression. Particularly notable is its performance in the 64-digit passkey retrieval task, where it substantially outperforms previously established methods such as H2O under the same compression conditions.

Implications and Future Developments

The implications of LagKV are twofold: practically, it offers a strategy that is easy to integrate into existing inference frameworks; theoretically, it challenges the reliance on attention weights for token importance determination. Given its performance and compatibility with FA, LagKV presents a promising pathway for enhancing the efficiency of LLMs without compromising on accuracy. As AI models strive for greater performance while being resource-conscious, future AI developments may explore hybrid strategies that combine elements of LagKV with other methods for even greater efficiency. Additionally, there may be scope for extending this approach through more nuanced considerations of token coherence and adaptive lag partitioning in the context of varied linguistic structures or additional modalities.

In conclusion, LagKV represents a significant advancement in the landscape of LLM deployment by effectively balancing computational efficiency and task accuracy. Its robust performance, free from attention mechanisms, sets a new precedent for future research in model optimization strategies.