[論文レビュー] Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time

Zichang Liu, Aditya Desai+6 more|arXiv (Cornell University)|May 26, 2023

Topic Modeling被引用数 10

ひとこと要約

本論文は Scissorhands を提案する。LLM 推論の KV キャッシュ圧縮法で、重要性仮説の持続性を利用して重要なトークンのみを保持し、ファインチューニング不要で KV キャッシュメモリを最大で 5×削減する。さらに 4 ビット量子化との互換性を示し、モデルの品質を維持する。

ABSTRACT

Large language models(LLMs) have sparked a new wave of exciting AI applications. Hosting these models at scale requires significant memory resources. One crucial memory bottleneck for the deployment stems from the context window. It is commonly recognized that model weights are memory hungry; however, the size of key-value embedding stored during the generation process (KV cache) can easily surpass the model size. The enormous size of the KV cache puts constraints on the inference batch size, which is crucial for high throughput inference workload. Inspired by an interesting observation of the attention scores, we hypothesize the persistence of importance: only pivotal tokens, which had a substantial influence at one step, will significantly influence future generations. Based on our empirical verification and theoretical analysis around this hypothesis, we propose Scissorhands, a system that maintains the memory usage of the KV cache at a fixed budget without finetuning the model. In essence, Scissorhands manages the KV cache by storing the pivotal tokens with a higher probability. We validate that Scissorhands reduces the inference memory usage of the KV cache by up to 5X without compromising model quality. We further demonstrate that Scissorhands can be combined with 4-bit quantization, traditionally used to compress model weights, to achieve up to 20X compression.

研究の動機と目的

• LLM 推論における KV キャッシュメモリのボトルネックを動機づけ、KV キャッシュメモリがモデルサイズを超え、バッチ処理量を制限する状況に対処する。
• ステップを通じて影響力を維持するトークンを特定するための Persistence of Importance Hypothesis を導入する。
• 再学習せずに精度を維持する予算意識的な KV キャッシュ圧縮法を開発する。
• 圧縮されたアテンションが元のアテンション出力を近似できるという理論的根拠を提供する。
• 圧縮をさらに高めるために量子化との互換性を示す。

提案手法

• 経験的に反復的なアテンションパターンを観察し、Persistence of Importance Hypothesis を定義する。
• 固定予算の下で高アテンション・トークンを保持し、影響力の薄いものを破棄して KV キャッシュを維持する Algorithm 1 および Algorithm 2 (Scissorhands) を開発する。
• 履歴ウィンドウを用いてトークンの重要性を推定し、貯水池サンプリングと LRU にインスパイアされた予算付き KV キャッシュ置換戦略を用いる。
• パワー-law 的なアテンションスコア分布の下で、元の KV キャッシュ出力と予算付き出力の期待近似誤差を示す理論的境界（Theorem 4.1）を提供する。
• Persistence decline に対処するために、後ろの層を優先してモデルの層とヘッドにメモリ予算を配分する。）

実験結果

リサーチクエスチョン

• RQ1 Persistence of Importance Hypothesis は異なる LLM および層に跨って成り立つか？
• RQ2 再学習なしで固定予算の下で KV キャッシュメモリを削減し、生成品質を維持できるか？
• RQ3 予算付き KV キャッシュがアテンション出力および下流タスクに及ぼす理論的・経験的影響は何か？
• RQ4 Scissorhands は 4-bit 量子化のようなポストトレーニング量子化と互換性があるか？

主な発見

• Scissorhands は OPT モデルにおいて、言語モデリングおよび下流タスクでモデル品質を劣化させることなく KV キャッシュ memory を最大で 5×削減する。
• Persistence ratio は層間で一般に 95% を超え、初期の重要トークンが後の重要性を大きく決定することを示している。
• Attention-score-based token dropping は元モデルと類似の出力を維持し、OPT-66B で最大 5×圧縮まで精度低下はほとんどない。
• 量子化（4-bit）は Scissorhands と互換性を維持し、評価タスクで性能を維持する。
• 理論分析（Theorem 3.1 および 4.1）は、アテンションウェイト、MLP スキップ接続、および予算付きトークン保持を近似挙動と誤差境界に結びつける。