TreeKV: Smooth Key-Value Cache Compression with Tree Structures

要約

効率的なキーバリュー (KV) キャッシュ圧縮は、長いシーケンスやリソースが制限された設定でトランスフォーマーベースの大規模言語モデル (LLM) をスケーリングするために重要です。
既存の方法では、位置または重要度スコアに基づいてトークンを排除しますが、位置ベースの戦略では、事前定義された領域外の重要な情報を見逃す可能性があります。一方、グローバル重要度スコアに依存する方法では、強い地域偏りが生じ、KV キャッシュ全体のコンテキスト保持が制限され、パフォーマンスが低下する可能性があります。
複雑なタスクに対する LLM の使用。
私たちのウェーブレット分析により、トークンがシーケンスの終わりに近づくにつれて、生成への寄与が徐々に増加し、隣接するトークンからさらに乖離する傾向があることが明らかになりました。これは、遠くのコンテキストから近くのコンテキストへの複雑さと変動性が増大して滑らかに移行していることを示しています。
この観察に動機付けられて、私たちは、スムーズなキャッシュ圧縮のためにツリー構造を採用した直感的でトレーニング不要の方法である TreeKV を提案します。
TreeKV は固定キャッシュ サイズを維持するため、LLM は長いテキスト シナリオでも高品質の出力を提供できます。
ほとんどの圧縮方法とは異なり、TreeKV は生成段階と事前入力段階の両方に適用できます。
これは、PG19 および OpenWebText2 の言語モデリング タスクにおけるすべてのベースライン モデルを常に上回っており、短いコンテキスト ウィンドウでトレーニングされた LLM を 16 倍のキャッシュ削減により長いウィンドウに一般化できるようになります。
Longbench ベンチマークでは、TreeKV は最適な効率で予算のわずか 6\% で最高のパフォーマンスを達成します。

要約(オリジナル)

Efficient key-value (KV) cache compression is critical for scaling transformer-based Large Language Models (LLMs) in long sequences and resource-limited settings. Existing methods evict tokens based on their positions or importance scores, but position-based strategies can miss crucial information outside predefined regions, while those relying on global importance scores resulting in strong regional biases, limiting the KV cache’s overall context retention and potentially impairing the performance of LLMs on complex tasks. Our wavelet analysis reveals that as tokens approach the end of sequence, their contributions to generation gradually increase and tends to diverge more from neighboring tokens, indicating a smooth transition with increasing complexity and variability from distant to nearby context. Motivated by this observation, we propose TreeKV, an intuitive, training-free method that employs a tree structure for smooth cache compression. TreeKV maintains a fixed cache size, allowing LLMs to deliver high-quality output even in long text scenarios. Unlike most compression methods, TreeKV is applicable to both the generation and prefilling stages. It consistently surpasses all baseline models in language modeling tasks on PG19 and OpenWebText2, allowing LLMs trained with short context window to generalize to longer window with a 16x cache reduction. On the Longbench benchmark, TreeKV achieves the best performance with only 6\% of the budget at optimal efficiency.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google