ThinK: Thinner Key Cache by Query-Driven Pruning

要約

大規模言語モデル (LLM) は自然言語処理の分野に革命をもたらし、増大したモデル サイズとシーケンス長を活用することで、さまざまなアプリケーションにわたって前例のないパフォーマンスを実現しました。
ただし、これに伴う計算コストとメモリ コストの増加は、特にトランスフォーマー アテンション メカニズムの二次的な複雑さによる長いシーケンスの管理において、重大な課題を引き起こします。
このペーパーでは、ロングコンテキストのシナリオに焦点を当て、推論中の KV キャッシュ メモリ消費の非効率性に対処します。
シーケンスの長さに基づいてメモリを最適化する既存のアプローチとは異なり、KV キャッシュのチャネル次元が、不均衡な大きさの分布とアテンションの重みの低ランク構造によって特徴づけられる、重大な冗長性を示すことを発見しました。
これらの観察に基づいて、最も重要でないチャネルを選択的にプルーニングしながら、アテンションの重み損失を最小限に抑えるように設計された、新しいクエリ依存の KV キャッシュ プルーニング手法である ThinK を提案します。
私たちのアプローチは、モデルの精度を維持または向上させるだけでなく、通常の KV キャッシュ削除方法と比較してメモリ コストの 20% 以上の削減も達成します。
さまざまな長期シーケンス データセットにわたる LLaMA3 モデルと Mistral モデルの広範な評価により、ThinK の有効性が確認され、パフォーマンスを損なうことなく効率的な LLM 導入の新たな前例が確立されました。
また、値キャッシュ プルーニングに手法を拡張する可能性についても概説し、メモリと計算のオーバーヘッドの両方を削減する際の ThinK の多用途性と幅広い適用性を示します。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google