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

要約

大規模言語モデル（Large Language Models: LLM）の急速な進歩により、現代のアプリケーションでは、拡張された文脈シーケンスを処理する需要が高まっている。しかし、この進歩は2つの大きな課題に直面している。すなわち、配列の長さが分布から外れることによる性能低下と、注意の2次計算複雑性によって引き起こされる推論時間の過剰な長さである。これらの問題は、LLMのロングコンテキストシナリオへの応用を妨げている。本論文では、効率的で高精度なロングコンテキスト推論のための、訓練不要の手法である動的トークンレベルKVキャッシュ選択（Dynamic Token-Level KV Cache Selection：TokenSelect）を提案する。TokenSelectは、非連続的なアテンションスパース性の観測に基づき、Query-Keyのドット積を用いて、トークン単位でヘッド毎のKVキャッシュのクリティカリティを測定する。ヘッドごとのソフトな投票メカニズムにより、TokenSelectは精度を犠牲にすることなく、少数のクリティカルなKVキャッシュ・トークンを選択的にアテンション計算に関与させる。TokenSelectをさらに高速化するために、連続するクエリの類似性の観測に基づいて選択キャッシュを設計し、効率的なドット積カーネルを実装することで、オーバーヘッドを大幅に削減した。TokenSelectの包括的な評価により、注目度計算において最大23.84倍の高速化、エンドツーエンドの待ち時間において最大2.28倍の高速化が実証され、同時に、最新のロングコンテキスト推論手法と比較して優れた性能を提供する。

要約(オリジナル)

The rapid advancement of Large Language Models (LLMs) has driven growing demand for processing extended context sequences in contemporary applications. 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 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 few critical KV cache tokens in attention calculation without sacrificing accuracy. To further accelerate TokenSelect, we design the Selection Cache based on observations of consecutive Query similarity and implemented efficient dot product kernel, significantly reducing the overhead. 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, DeepL