[論文レビュー] H$_2$O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models

Zhenyu Zhang, Ying Sheng+10 more|arXiv (Cornell University)|Jun 24, 2023

Topic Modeling被引用数 27

ひとこと要約

H2Oは、ヘビー・ヒッターに基づく KV-cache のエビクションポリシーを導入し、メモリ使用量を大幅に削減し、LLM推論のスループットを向上させつつ精度を損なうことなく、OPT、LLaMA、GPT-NeoXで検証済み。

ABSTRACT

Hyde-IKV is a dynamic management system for the Key-Value (KV) cache in Small Language Models (SLMs). It tackles the memory bottleneck of long-context inference by intelligently prioritizing and compressing the stored context. Instead of abruptly deleting tokens, Hyde-IKV uses a hybrid strategy of positional decay (gradually reducing the influence of older tokens) and importance-based modulation (preserving tokens with high attention scores or large key vector magnitudes). This non-invasive approach, implemented via hooks into the transformer architecture, significantly reduces GPU memory usage and speeds up inference for long sequences while largely preserving the quality of the generated text.

研究の動機と目的

• 長い内容の生成における KV-cache のメモリ削減を、精度の大幅な低下を伴わずに動機づける。
• 少数のトークン（ヘビー・ヒッター）が大半のアテンションスコアを支配していることを特定する。
• ヘビー・ヒッターを活用して最近性と関連性のバランスを取る、低コストのエビクションポリシーを開発する。
• 動的サブモジュラエビクションフレームワークの理論的保証を提供する。
• 複数のモデルファミリでスループットとレイテンシの実際的な改善を示す。

提案手法

• アテンションのスパース性を観察（95％超）し、より小さな KV-cache の必要性を正当化する。
• ヘビー・ヒッター（H2）を、 Concentrated attention scoresと頻繁な語と高い共起を持つトークンとして特徴づける。
• KV-cache のエビクションを動的サブモジュラ問題として定式化し、F_score に基づく H2O エビクション（グリーディ法）を提案する。
• デコードごとの局所アテンション統計を用いる、低コストで動的なグリーディアルゴリズムを実装する。
• 緩やかなサブモジュラ性の仮定の下で、グリーディポリシーのほぼ最適性保証を証明する。
• FlexGen への H2O の統合により、OPT、LLaMA、GPT-NeoX の 20% KV-cache 予算で評価する。

実験結果

リサーチクエスチョン

• RQ1 KV-cache のメモリを大幅に削減できるか、生成品質を損なわずに？
• RQ2 ヘビー・ヒッター・トークン（H2）は大半のアテンション信号を捉え、エビクション判断を導けるか？
• RQ3 グリーディで低コストのエビクションポリシーが、サブモジュラなアテンション仮定の下でほぼ最適な性能を達成できるか？
• RQ4 H2Oは、モデルファミリ（OPT、LLaMA、GPT-NeoX）およびシーケンス長が異なるタスクでどのように性能を発揮するか？
• RQ5 H2Oを実 systems で使用した場合、どの程度のスループットとレイテンシの改善が実現可能か？

主な発見

• KV-cache を約20%まで削減すると、多くのタスクで同等またはそれ以上の生成品質を得られ、メモリ節約も大幅。
• H2Oは OPT-6.7B および OPT-30B で DeepSpeed Zero-Inference および Hugging Face Accelerate に対して最大で 29×、FlexGen に対して最大で 3×のスループット向上を達成。
• H2Oは同じバッチサイズでベースラインシステムと比較してレイテンシを最大で 1.9×低減。
• 既存のスパース化手法と H2 を組み合わせることで、低予算で他の手法が失敗する場面でも性能を維持。
• このアプローチは、モデルサイズ（6.7B–175B）および HELM や lm-eval-harness のタスクを含む長いシーケンス生成を含め、堅牢性を示す。