ThinK: Thinner Key Cache by Query-Driven Pruning
Yuhui Xu, Zhanming Jie, Hanze Dong, Lei Wang, Xudong Lu, Aojun Zhou, Amrita Saha, Caiming Xiong, Doyen Sahoo
2024-07-31
Summary
This paper presents ThinK, a new method designed to reduce the memory usage of key-value (KV) caches in large language models (LLMs) during processing. It focuses on improving efficiency while maintaining or enhancing model performance.
What's the problem?
As large language models become more complex and handle longer sequences of data, they require more memory and computational power. This can lead to inefficiencies, especially when using KV caches, which store important information for quick access. Traditional methods for optimizing memory often do not address the specific needs of these models, resulting in wasted resources and slower performance.
What's the solution?
To tackle this issue, the authors developed ThinK, which uses a method called query-driven pruning. This approach identifies and removes the least important channels in the KV cache based on their relevance to the current task. By focusing on which parts of the cache are necessary for each specific query, ThinK reduces memory costs by over 20% compared to standard methods while still keeping model accuracy high. The authors tested ThinK on various models and found it significantly improved efficiency without sacrificing performance.
Why it matters?
This research is important because it offers a practical solution to one of the major challenges in deploying large language models, particularly in environments with limited resources. By making these models more efficient, ThinK can help improve their use in real-world applications like chatbots, translation services, and other AI-driven technologies that require fast and accurate processing of long sequences.
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.