Compressing KV Cache for Long-Context LLM Inference with Inter-Layer Attention Similarity (2412.02252v1)

Abstract: The increasing context window size in LLMs, such as the GPT and LLaMA series, has improved their ability to tackle complex, long-text tasks, but at the cost of inference efficiency, particularly regarding memory and computational complexity. Existing methods, including selective token retention and window-based attention, improve efficiency but risk discarding important tokens needed for future text generation. In this paper, we propose an approach that enhances LLM efficiency without token loss by reducing the memory and computational load of less important tokens, rather than discarding them.We address two challenges: 1) investigating the distribution of important tokens in the context, discovering recent tokens are more important than distant tokens in context, and 2) optimizing resources for distant tokens by sharing attention scores across layers. The experiments show that our method saves $35\%$ KV cache without compromising the performance.

• The paper introduces POD, a novel approach that compresses the KV cache by focusing on proximal tokens, maintaining 80% next-token prediction accuracy.
• It leverages inter-layer attention similarity by sharing attention scores across layers, which reduces computational overhead and increases maximum batch size by 30%.
• The method significantly improves inference efficiency in long-context LLMs while preserving performance, opening avenues for further research on architectural optimizations.

An Efficient Approach to KV Cache Compression in LLM Inference

The expansion of context window sizes in LLMs, such as the GPT and LLaMA series, has created significant challenges in managing memory and computational complexity particularly during inference. Such challenges arise because the computational complexity of the attention mechanism increases quadratically with the context window size, while the size of the KV cache grows linearly. Current strategies aimed at improving inference efficiency often involve discarding tokens, which risks losing valuable information needed for subsequent text generation.

In tackling the issues of heightened inference inefficiency without sacrificing performance, the work discussed here introduces an innovative approach termed Proximal tokens over Distant tokens (POD). POD strategically compresses the KV cache by focusing computational and memory resources on the most impactful tokens while sharing attention scores for less critical distant tokens across layers.

Insights and Methodology

The core insights driving the POD methodology are twofold:

1. Token Importance Distribution: The analysis highlights that proximal tokens (comprising both initial and recently added tokens) have a higher significance compared to distant tokens. Experiments indicate that models achieve identical next-token predictions by attending only to proximal tokens in 80% of scenarios compared to attending to all tokens.
2. Inter-Layer Attention Similarity: The paper reveals that for consecutive layers within LLMs, attention scores manifest similar patterns—a phenomenon scaled to modern LLMs and utilized in POD to share attention scores for distant tokens exclusively.

Based on these observations, the methodology is divided into three stages:

• Exploration of Offline Inter-Layer Attention Sharing: Groups layers based on shared attention scores.
• Lightweight Training Adaptation: Refines models by post-training them to recognize and utilize the shared attention patterns.
• Efficient Inference: Implements attention sharing for distal tokens, allowing models to optimize memory usage and computation by only caching crucial proximal token states.

Experimental Results and Impact

Experiments conducted on benchmark datasets demonstrate that POD can effectively reduce KV cache size by 35% without degrading model performance. This reduction is substantial compared to previous methods, with empirical evaluations showing a 30% increase in maximum batch processing size. The probing of task-specific scenarios, such as long-context tasks on benchmarks like LongBench and LEval, affirms the utility of the proposed method. Notably, POD succeeds in maintaining the integrity of token streams where traditional methods might fail due to restricted token windows.

Implications and Future Directions

The implications of this research span both practical and theoretical realms. Practically, the method stands to significantly decrease computational overheads in large-scale deployment scenarios. Theoretically, the findings encourage further exploration into the utilization of inter-layer redundancies for various architectural optimizations beyond KV cache reduction.

Future enhancements to the POD methodology could involve exploring the adaptive tuning of proximal and distant token categorization and extending the model to support even longer context windows seen in nascent ultra-long context LLMs. Additionally, integrating POD with other resource-reduction techniques like token quantization may further improve efficiency while preserving performance in increasingly complex NLP tasks.

As the demand for LLMs evolves, methodologies such as POD offer promising avenues for sustaining high-performance inference amid growing model scale and complexity.