ThinK: Thinner Key Cache by Query-Driven Pruning

Original Paper: https://arxiv.org/abs/2407.21018

By: Yuhui Xu, Zhanming Jie, Hanze Dong, Lei Wang, Xudong Lu, Aojun Zhou, Amrita Saha, Caiming Xiong, Doyen Sahoo

Abstract:

Large Language Models (LLMs) have revolutionized the field of natural language processing, achieving unprecedented performance across a variety of applications by leveraging increased model sizes and sequence lengths. However, the associated rise in computational and memory costs poses significant challenges, particularly in managing long sequences due to the quadratic complexity of the transformer attention mechanism. This paper focuses on the long-context scenario, addressing the inefficiencies in KV cache memory consumption during inference. Unlike existing approaches that optimize the memory based on the sequence lengths, we uncover that the channel dimension of the KV cache exhibits significant redundancy, characterized by unbalanced magnitude distribution and low-rank structure in attention weights. Based on these observations, we propose ThinK, a novel query-dependent KV cache pruning method designed to minimize attention weight loss while selectively pruning the least significant channels. Our approach not only maintains or enhances model accuracy but also achieves a reduction in memory costs by over 20% compared with vanilla KV cache eviction methods. Extensive evaluations on the LLaMA3 and Mistral models across various long-sequence datasets confirm the efficacy of ThinK, setting a new precedent for efficient LLM deployment without compromising performance. We also outline the potential of extending our method to value cache pruning, demonstrating ThinK's versatility and broad applicability in reducing both memory and computational overheads.

Summary Notes

Figure 1:An illustration of the pruning procedure of ThinK. Within each head, scores are calculated for each channel, and only the top T channels out of D are selected for retention. A binary channel mask, along with the pruned keys, is subsequently stored in the cache memory.

In recent years, Large Language Models (LLMs) have been at the forefront of advances in natural language processing, enabling impressive feats in a variety of applications such as document summarization, code generation, and conversational AI. However, the computational and memory costs of managing these models, especially when dealing with long sequences, have been a significant barrier to their deployment. The transformer architecture, which underpins these models, suffers from quadratic complexity in its attention mechanism, making it particularly challenging to manage long contexts efficiently.

In this blog post, we delve into a groundbreaking approach to this problem, as presented in the research paper titled "ThinK: Thinner Key Cache by Query-Driven Pruning." This innovative method promises to significantly reduce memory and computational overheads without compromising model accuracy.

Introduction: The Challenge of Long Sequences

The core of the issue lies in the key-value (KV) cache used by transformers during inference. The KV cache size is a product of multiple factors including batch size, sequence length, number of layers, number of heads, and channel size. As these models scale, managing this cache becomes increasingly burdensome. Traditionally, methods like quantization and token eviction have been used to alleviate this issue, but they often overlook the redundancy present in the channel dimension of the KV cache.

Methodology: The ThinK Approach

ThinK introduces a novel approach to KV cache pruning that leverages the redundancy in the channel dimension. The key insight from the research is that the magnitude distribution across the KV cache's channels is significantly unbalanced, and the attention weights exhibit a low-rank structure. This suggests that many channels contribute little to the model's performance and can be pruned.|

Query-Driven Pruning

The ThinK method formulates the pruning task as an optimization problem, aiming to minimize the loss in attention weights due to pruning. The approach involves a query-dependent criterion that assesses the importance of each channel. By calculating scores for each channel based on their interaction with the query vectors, ThinK identifies and retains only the most critical channels.

Here's a high-level overview of the process:

1. Score Calculation: For each channel in the key cache, calculate the importance score using the query and key vectors.
2. Greedy Selection: Retain the top channels with the highest scores.
3. Pruning: Prune the less significant channels, reducing the overall size of the KV cache.

This method ensures that the essential information required for accurate attention computation is preserved, while redundant data is discarded.

Implementation Details

Two implementations of ThinK during the decoding stage are proposed:

1. Zero-Filling Pruned Keys: The pruned keys are zero-filled to restore their original size before concatenation with unpruned keys.
2. Pruned Query Multiplication: The pruned query is multiplied by the pruned key, and the unpruned query is applied to the unpruned key. The results are then concatenated.

Both implementations are designed to be integrated seamlessly with other optimization techniques like FlashAttention.

Experimental Results

The efficacy of ThinK was validated using the LLaMA3 and Mistral models across various long-sequence datasets from benchmarks such as LongBench and Needle-in-a-Haystack. The results were impressive:

• Memory Reduction: ThinK achieved a reduction in KV cache memory costs by more than 20%.
• Maintained Accuracy: The method maintained or even enhanced model accuracy compared to traditional KV cache eviction methods.
• Robustness: ThinK was able to maintain performance across varying KV cache sizes and pruning ratios.

LongBench Results

Tables from the research paper showed that ThinK not only reduced memory usage but also improved performance in many cases. For instance, with a KV cache size of 2048 and a pruning ratio of 40%, ThinK outperformed the LLaMA3-8B model with a full KV cache.

Needle-in-a-Haystack Results

ThinK demonstrated robust performance in the Needle-in-a-Haystack test, particularly with larger KV cache sizes, where it maintained or improved accuracy compared to the baseline SnapKV method.

Conclusion and Future Directions

ThinK represents a significant advancement in the efficient management of KV caches in LLMs. By focusing on the channel dimension and employing a query-driven pruning strategy, ThinK reduces memory and computational overheads while maintaining model performance. This method is orthogonal to existing KV cache compression schemes, making it a versatile tool for enhancing LLM efficiency.

Future research will aim to increase the pruning ratio without compromising performance and explore pruning techniques for value caches. Additionally, more sophisticated methods to assess channel importance, incorporating both token-level and channel-level information, could further refine this approach.

In summary, ThinK offers a promising solution to one of the major challenges in deploying large-scale language models, paving the way for more efficient and practical applications of these powerful tools.