ThinK: クエリ駆動型プルーニングによるキーキャッシュの軽量化

ThinK: Thinner Key Cache by Query-Driven Pruning

July 30, 2024

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

cs.AI

要旨

大規模言語モデル（LLMs）は、モデルサイズとシーケンス長の増大を活用することで、自然言語処理の分野に革命をもたらし、様々なアプリケーションで前例のない性能を達成してきました。しかし、それに伴う計算コストとメモリコストの増大は、特にトランスフォーマーのアテンションメカニズムの二次複雑性による長いシーケンスの管理において、重大な課題を提起しています。本論文は、長文脈シナリオに焦点を当て、推論時のKVキャッシュメモリ消費の非効率性に対処します。既存のアプローチがシーケンス長に基づいてメモリを最適化するのとは異なり、我々はKVキャッシュのチャネル次元に大きな冗長性があることを明らかにしました。これは、アテンション重みの不均衡な大きさ分布と低ランク構造によって特徴づけられます。これらの観察に基づき、我々はThinKという新しいクエリ依存型KVキャッシュプルーニング手法を提案します。この手法は、アテンション重みの損失を最小化しつつ、最も重要でないチャネルを選択的にプルーニングするように設計されています。我々のアプローチは、モデルの精度を維持または向上させるだけでなく、従来のKVキャッシュ削除方法と比較してメモリコストを20%以上削減します。LLaMA3とMistralモデルを用いた様々な長文データセットでの広範な評価により、ThinKの有効性が確認され、性能を損なうことなく効率的なLLM展開の新たな先例を確立しました。また、我々の手法をバリューキャッシュプルーニングに拡張する可能性についても概説し、ThinKがメモリと計算オーバーヘッドの両方を削減するための汎用性と幅広い適用性を実証しました。

English

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.