TidalDecode: 位置持続スパース注意を用いた高速かつ正確なLLMデコーディング

TidalDecode: Fast and Accurate LLM Decoding with Position Persistent Sparse Attention

October 7, 2024

著者: Lijie Yang, Zhihao Zhang, Zhuofu Chen, Zikun Li, Zhihao Jia

cs.AI

要旨

大規模言語モデル（LLMs）は、さまざまな自然言語処理タスクで重要な進展を遂げてきました。長いコンテキストを扱うために注目されているモデルが増えています。ただし、Transformerアーキテクチャに必要な拡張されたキー・バリュー（KV）キャッシュサイズの増加は、デコーディングフェーズ中に特にメモリ制約を強化し、大きなボトルネックを作り出しています。このボトルネックに対処するために設計された既存の疎な注意機構には、2つの制限があります。まず、最も関連性の高いトークンを信頼性を持って特定することができないこと、そして連続するTransformerレイヤー間でのトークン選択の空間的な一貫性を見落としていることが挙げられます。この論文では、位置持続型の疎な注意を通じて高速かつ正確なLLMデコーディングを実現するための、シンプルで効果的なアルゴリズムおよびシステムであるTidalDecodeを紹介します。TidalDecodeは、既存の疎な注意方法によって選択されたトークンの空間的な一貫性を活用し、最高の注意スコアを持つトークンを特定するためのいくつかのトークン選択レイヤーを導入し、他のすべてのレイヤーが事前に選択されたトークンを使用して疎な注意を行うようにします。この設計により、TidalDecodeは、疎な注意のトークン選択のオーバーヘッドを大幅に削減すると同時に、生成された結果の品質を犠牲にすることなく、LLMデコーディングの遅延を最大2.1倍まで削減できます。様々なLLMsおよびタスクに対する評価により、TidalDecodeが完全な注意方法の生成パフォーマンスに密接に匹敵する一方、LLMデコーディングの遅延を最大2.1倍まで削減することが示されました。

English

Large language models (LLMs) have driven significant advancements across diverse NLP tasks, with long-context models gaining prominence for handling extended inputs. However, the expanding key-value (KV) cache size required by Transformer architectures intensifies the memory constraints, particularly during the decoding phase, creating a significant bottleneck. Existing sparse attention mechanisms designed to address this bottleneck have two limitations: (1) they often fail to reliably identify the most relevant tokens for attention, and (2) they overlook the spatial coherence of token selection across consecutive Transformer layers, which can lead to performance degradation and substantial overhead in token selection. This paper introduces TidalDecode, a simple yet effective algorithm and system for fast and accurate LLM decoding through position persistent sparse attention. TidalDecode leverages the spatial coherence of tokens selected by existing sparse attention methods and introduces a few token selection layers that perform full attention to identify the tokens with the highest attention scores, while all other layers perform sparse attention with the pre-selected tokens. This design enables TidalDecode to substantially reduce the overhead of token selection for sparse attention without sacrificing the quality of the generated results. Evaluation on a diverse set of LLMs and tasks shows that TidalDecode closely matches the generative performance of full attention methods while reducing the LLM decoding latency by up to 2.1x.