Efficient Sparse Attention needs Adaptive Token Release

要約

近年、ラージ言語モデル (LLM) は、テキスト中心のさまざまなタスクにわたって優れた機能を実証してきました。
ただし、その「大規模な」スケールにより、特にトランスフォーマーのキーと値の状態の管理において、計算とストレージに重大な課題が生じ、その広範な適用性が制限されます。
したがって、キャッシュからリソースを適応的に解放し、必要なキーと値の状態を再構築することを提案します。
特に、軽量のコントローラー モジュールによってこれを実現し、理想的なトップ千ドルのまばらな注意力に近づけます。
このモジュールは、最も高い $K$ アテンションの重みを持つトークンを保持し、同時に破棄されたものの必要なトークンを再構築します。これは将来のデコードに不可欠になる可能性があります。
自然言語の生成とモデリングにおける包括的な実験により、私たちの方法がパフォーマンスの点で十分な注意を払って競争力があるだけでなく、最大 221.8% という大幅なスループットの向上を達成していることが明らかになりました。
レプリケーションのコードは https://github.com/WHUIR/ADORE で入手できます。

要約(オリジナル)

In recent years, Large Language Models (LLMs) have demonstrated remarkable capabilities across a wide array of text-centric tasks. However, their `large’ scale introduces significant computational and storage challenges, particularly in managing the key-value states of the transformer, which limits their wider applicability. Therefore, we propose to adaptively release resources from caches and rebuild the necessary key-value states. Particularly, we accomplish this by a lightweight controller module to approximate an ideal top-$K$ sparse attention. This module retains the tokens with the highest top-$K$ attention weights and simultaneously rebuilds the discarded but necessary tokens, which may become essential for future decoding. Comprehensive experiments in natural language generation and modeling reveal that our method is not only competitive with full attention in terms of performance but also achieves a significant throughput improvement of up to 221.8%. The code for replication is available on the https://github.com/WHUIR/ADORE.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google