EMS: Adaptive Evict-then-Merge Strategy for Head-wise KV Cache Compression Based on Global-Local Importance

要約

大規模言語モデル (LLM) が進歩し続けるにつれて、さまざまなアプリケーションにわたる長いコンテキストの高品質かつ高速な処理に対する需要が高まっています。
KV キャッシュは、以前に生成されたキーと値のトークンを保存し、推論中の冗長な計算を効果的に削減するため、広く採用されています。
しかし、メモリのオーバーヘッドが重大な問題となるにつれ、KV キャッシュの効率的な圧縮に対する注目が高まっています。
既存の方法のほとんどは、重要なトークンの特定と圧縮戦略の設計という 2 つの観点から圧縮を実行します。
ただし、これらのアプローチでは、蓄積された注意スコアや位置エンコーディングの影響により、重要なトークンの偏った分布が生成されることがよくあります。
さらに、異なるヘッド間の疎性と冗長性を見落としているため、ヘッドレベルで最も有効な情報を保存することが困難になります。
この目的を達成するために、極端な圧縮率の下でより優れた KV キャッシュ圧縮を達成しながら、これらの制限を克服する EMS を提案します。
具体的には、トークンの重要性をより適切に識別するために、グローバル KV トークンとローカル KV トークンの両方から蓄積されたアテンション スコアを組み合わせたグローバル-ローカル スコアを導入します。
圧縮戦略については、さまざまなヘッドにわたる KV トークンの疎性と冗長性を考慮した、適応性のある統合された Evict-then-Merge フレームワークを設計します。
さらに、効率を高めるために、ゼロクラス メカニズムを通じてヘッド単位の並列圧縮を実装します。
広範な実験により、極端な圧縮率下でも SOTA パフォーマンスが実証されました。
EMS は一貫して最低のパープレキシティを達成し、256 のキャッシュ バジェットで LongBench 上の 4 つの LLM にわたってスコアを 1.28 ポイント以上改善し、ニードルイン ア コンテキスト長の 2% 未満のキャッシュ バジェットで 95% の取得精度を維持します。
干し草の山のタスク。

要約(オリジナル)

As large language models (LLMs) continue to advance, the demand for higher quality and faster processing of long contexts across various applications is growing. KV cache is widely adopted as it stores previously generated key and value tokens, effectively reducing redundant computations during inference. However, as memory overhead becomes a significant concern, efficient compression of KV cache has gained increasing attention. Most existing methods perform compression from two perspectives: identifying important tokens and designing compression strategies. However, these approaches often produce biased distributions of important tokens due to the influence of accumulated attention scores or positional encoding. Furthermore, they overlook the sparsity and redundancy across different heads, which leads to difficulties in preserving the most effective information at the head level. To this end, we propose EMS to overcome these limitations, while achieving better KV cache compression under extreme compression ratios. Specifically, we introduce a Global-Local score that combines accumulated attention scores from both global and local KV tokens to better identify the token importance. For the compression strategy, we design an adaptive and unified Evict-then-Merge framework that accounts for the sparsity and redundancy of KV tokens across different heads. Additionally, we implement the head-wise parallel compression through a zero-class mechanism to enhance efficiency. Extensive experiments demonstrate our SOTA performance even under extreme compression ratios. EMS consistently achieves the lowest perplexity, improves scores by over 1.28 points across four LLMs on LongBench under a 256 cache budget, and preserves 95% retrieval accuracy with a cache budget less than 2% of the context length in the Needle-in-a-Haystack task.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google