Locret: Enhancing Eviction in Long-Context LLM Inference with Trained Retaining Heads on Consumer-Grade Devices

要約

大規模な言語モデル（LLM）の入力コンテキストの長さをスケーリングすると、計算コストとメモリフットプリントが大幅に増加し、注意キー価値（kV）キャッシュを維持します。
既存のKVキャッシュ圧縮法は、非効率的な圧縮戦略と限られたメモリ削減効果に悩まされているため、特に長いコンテキストストリーム入力を推測する場合、LLMが消費者グレードのデバイスに長いコンテキスト推論を行うことが困難です。
このような障害は、消費者グレードのデバイスがより複雑なアプリケーションをサポートすることを妨げ、LLMの民主化に対する課題を生み出します。
これを克服するために、Chunked Prefillと互換性のある立ち退きポリシーを作成する最初のフレームワークであるLocretを提案します。
Locretは、学習可能な保持ヘッドによってKVキャッシュユニットの因果関係の重要性を評価することにより、キャッシュユニットの正確な立ち退きを可能にし、効率的な長いコンテキスト推論を促進します。
私たちの広範な経験的研究では、Locretは、メモリ効率と生成品質の点で最近の一般的で競争力のあるアプローチを上回ります。Locretは、パフォーマンスの10％未満で最大20倍のKVキャッシュ圧縮率を達成します。
さらに、Locretは、生成品質を損なうことなく、単一のNVIDIA 4090 GPUで128K+長いコンテキスト推論を達成し、追加のトレーニングの1時間未満のコストのみを費用します。

要約(オリジナル)

Scaling the input context length of a large language model (LLM) incurs a significant increase in computation cost and memory footprint to maintain the attention key-value (KV) cache. Existing KV cache compression methods suffer from inefficient compression strategies and limited memory reduction effects, making it difficult for LLMs to conduct long-context inference on consumer-grade devices, especially when inferring long-context stream input. Such obstacles prevent consumer-grade devices from supporting more complex applications, creating challenges for the democratization of LLMs. To overcome this, we propose Locret, the first framework to create an eviction policy compatible with chunked prefill. By evaluating the causal importance of KV cache units by learnable retaining heads, Locret enables precise eviction of cache units, facilitating efficient long-context inference. In our extensive empirical studies, Locret outperforms the recent popular and competitive approaches in terms of memory efficiency and generation quality — Locret achieves up to 20x of KV cache compression ratio within less than 10% performance loss. Furthermore, Locret achieves 128K+ long-context inference on a single NVIDIA 4090 GPU without compromising generation quality and only costs <1 GPU hour of additional training.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google