[論文レビュー] Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs

Suyu Ge, Yunan Zhang+4 more|arXiv (Cornell University)|Oct 3, 2023

Topic Modeling被引用数 5

ひとこと要約

FastGenは、軽量なアテンションプロファイリングを活用して、動的に適応的な圧縮ポリシーを適用することで、LLMにおけるKVキャッシュの圧縮を実現する。具体的には、長距離コンテキストの削除、特別トークンの保持、または低頻度トークンの破棄といった、最適化された圧縮ポリシーを動的に適用する。その結果、微細調整や再トレーニングを必要とせず、生成品質の損失を最小限に抑えながら最大36%のメモリ削減を達成する。

ABSTRACT

In this study, we introduce adaptive KV cache compression, a plug-and-play method that reduces the memory footprint of generative inference for Large Language Models (LLMs). Different from the conventional KV cache that retains key and value vectors for all context tokens, we conduct targeted profiling to discern the intrinsic structure of attention modules. Based on the recognized structure, we then construct the KV cache in an adaptive manner: evicting long-range contexts on attention heads emphasizing local contexts, discarding non-special tokens on attention heads centered on special tokens, and only employing the standard KV cache for attention heads that broadly attend to all tokens. Moreover, with the lightweight attention profiling used to guide the construction of the adaptive KV cache, FastGen can be deployed without resource-intensive fine-tuning or re-training. In our experiments across various asks, FastGen demonstrates substantial reduction on GPU memory consumption with negligible generation quality loss. We will release our code and the compatible CUDA kernel for reproducibility.

研究の動機と目的

• 微細調整や再トレーニングを必要とせずに、自己回帰的LLM推論におけるKVキャッシュのメモリ容量を削減すること。
• KVキャッシュの効率的かつ標的的な圧縮を可能にする、アテンションモジュール内の構造的パターンを同定・活用すること。
• 各アテンションヘッドの機能的役割に応じて圧縮戦略を適応的に変更できる、即挿し可能な手法を開発すること。
• 構造に配慮したスマートなプルーニングにより、生成品質を保持したまま顕著なメモリ節約を達成すること。
• GPUメモリ圧力を軽減することで、大規模言語モデルの効率的なデバイス内デプロイを可能にすること。

提案手法

• 本手法は、アテンションマップを分析し、各ヘッドを5つの構造的カテゴリに分類する軽量なアテンションプロファイリングステップを採用する。
• プロファイリングの結果に基づき、5つの圧縮ポリシーのいずれかを適用する：局所的コンテキストの削除、特別トークンの保持、頻度ベースのプルーニング、列方向のスパarsity化、または通常のキャッシュ。
• 周囲の頻度、特別トークン、標点、および局所的アテンションパターンに基づくポリシー集合に従い、関連するトークンのみを選択して、適応的KVキャッシュを構築する。
• プロファイリングのオーバーヘッドは非常に小さく、KVキャッシュメモリのたった0.78%にとどまるため、リアルタイム推論に実用的である。
• 圧縮ポリシーを適用する順序を、特別トークンや頻出トークンといった高効果のものから優先順位をつけるグリーディーなポリシー順序戦略を採用する。
• 従来のLLM推論パイプラインと完全に互換性があり、モデルの再トレーニングやアーキテクチャの変更を一切必要としない。

Figure 1: Different attention heads usually have different structures. Left: Four common attention structures (more details are elaborated in Section 3 and Section 4 ). Right: Attention map compositions of three attention heads that are in the same layer.

実験結果

リサーチクエスチョン

• RQ1 LLMのアテンションモジュールは、効率的かつ標的的なKVキャッシュ圧縮を可能にする、明確な構造的パターンに分類可能か？
• RQ2 生成品質を損なわず、各アテンションヘッドごとに適応的にKVキャッシュを圧縮する方法は何か？
• RQ3 メモリ節約を最大化しながらモデル性能を維持するための、圧縮ポリシーの最適な順序と組み合わせは何か？
• RQ4 即挿し可能で微細調整なしの手法が、KVキャッシュにおいて顕著なメモリ削減を達成し、品質損失をほとんど生じさせないか？
• RQ5 周波数や局所的アテンションのしきい値といったハイパーパrameterの選択に、圧縮性能はどれほど敏感か？

主な発見

• FastGenは、Llama 1-65BでKVキャッシュメモリ使用量を最大36.04%削減しながら、AlpacaEvalで49.75%の勝率を維持し、フルキャッシュベースラインとほぼ同等の性能を達成した。
• 特別トークンポリシー（${\bm{C}}_{\text{special}}$）を削除すると、性能が最も大きく低下し、2.11%の勝率低下が生じた。これは、そのポリシーが極めて重要な役割を果たしていることを示している。
• 頻度ベースのポリシー（${\bm{C}}_{\text{frequent}}$）が圧縮に最も貢献し、他のポリシーと組み合わせた場合、キャッシュを36.04%まで削減した。
• ポリシー適用順序が性能に顕著な影響を与える：デフォルトのグリーディー順序（特別トークン → 標点 → 頻出トークン → 局所的）が、最高の勝率と圧縮比を達成した。
• 本手法はハイパーパrameterの変更に対して頑健であり、全テスト設定で45%以上の勝率を維持した。一方、圧縮比は設定によって顕著に変動した。
• プロファイリングのオーバーヘッドは無視できるほど小さく、フルKVキャッシュサイズのたった0.78%にとどまり、リアルタイムデプロイに実用的である。

Figure 2: Performance of Adaptive KV Cache (FastGen) and Fixed KV Cache (Frequency, Local, and Frequency+Local; Zhang et al. , 2023 and Liu et al. , 2023a ) on AlpacaEval.