TokenSelect: Efficient Long-Context Inference and Length Extrapolation for LLMs via Dynamic Token-Level KV Cache Selection

要約

大規模言語モデル (LLM) の開発により、より長いコンテキストを処理する機能が、ドキュメント間の理解や LLM を利用した検索システムなどの Web アプリケーションにとって重要な機能になりました。
しかし、この進歩は 2 つの大きな課題に直面しています。1 つは分布外のシーケンス長によるパフォーマンスの低下、もう 1 つはアテンションの 2 次計算の複雑さによって引き起こされる過度に長い推論時間です。
これらの問題は、長いコンテキストのシナリオでの LLM の適用を妨げます。
この論文では、効率的かつ正確なロングコンテキスト推論のためのモデルに依存しない、トレーニング不要の方法である動的トークンレベル KV キャッシュ選択 (TokenSelect) を提案します。
TokenSelect は、非連続アテンション スパース性の観察に基づいて構築されており、Query-Key ドット積を使用してヘッドごとの KV キャッシュの重要性をトークン レベルで測定します。
ヘッドごとのソフト投票メカニズムにより、TokenSelect は、精度を犠牲にすることなく、少数の重要な KV キャッシュ トークンをアテンション計算に選択的に組み込みます。
TokenSelect をさらに高速化するために、連続するクエリの類似性の観察に基づいて選択キャッシュを設計し、効率的なドット積カーネルを実装して、トークン選択のオーバーヘッドを大幅に削減しました。
TokenSelect の包括的な評価では、最先端のロングコンテキスト推論手法と比較して優れたパフォーマンスを提供しながら、アテンション計算で最大 23.84 倍の高速化とエンドツーエンドのレイテンシで最大 2.28 倍の高速化が実証されました。

要約(オリジナル)

With the development of large language models (LLMs), the ability to handle longer contexts has become a key capability for Web applications such as cross-document understanding and LLM-powered search systems. However, this progress faces two major challenges: performance degradation due to sequence lengths out-of-distribution, and excessively long inference times caused by the quadratic computational complexity of attention. These issues hinder the application of LLMs in long-context scenarios. In this paper, we propose Dynamic Token-Level KV Cache Selection (TokenSelect), a model-agnostic, training-free method for efficient and accurate long-context inference. TokenSelect builds upon the observation of non-contiguous attention sparsity, using Query-Key dot products to measure per-head KV Cache criticality at token-level. By per-head soft voting mechanism, TokenSelect selectively involves a small number of critical KV cache tokens in the attention calculation without sacrificing accuracy. To further accelerate TokenSelect, we designed the Selection Cache based on observations of consecutive Query similarity and implemented efficient dot product kernel, significantly reducing the overhead of token selection. A comprehensive evaluation of TokenSelect demonstrates up to 23.84x speedup in attention computation and up to 2.28x acceleration in end-to-end latency, while providing superior performance compared to state-of-the-art long-context inference methods.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google