TokenButler: Token Importance is Predictable

要約

大規模な言語モデル（LLMS）は、キー価値（kV）キャッシュに依存してトークン履歴を保存し、トークンの効率的なデコードを可能にします。
KVキャッシュが成長するにつれて、それは主要なメモリと計算のボトルネックになりますが、特に以前の研究では、各デコードステップに有意義に貢献するトークンの小さなサブセットのみが示されているため、このボトルネックを緩和する機会があります。
これらの重要なトークンを見つける上での重要な課題は、それらが動的であり、非常に入力クエリに依存していることです。
既存の方法は、トークンを永続的に排除することで品質をリスクするか、完全なKVキャッシュを保持しますが、密集したコンテキストが豊富なタスクで失敗し、世代のトークンのチャンク（ページ）を取得することに依存しています。
さらに、多くの既存のKVキャッシュスパース性方法は、トークンの重要性のために不正確なプロキシに依存しています。
これらの制限に対処するために、これらの重要なトークンを特定することを学ぶ高粒度のクエリ認識予測子であるTokenbutlerを紹介します。
1.2％未満のパラメーターオーバーヘッドで軽量の予測因子をトレーニングすることにより、Tokenbutlerは、コンテキストの予測された重要性に基づいてトークンに優先順位を付けます。
これにより、トークンの重要性を推定するためのSOTAメソッドと比較して、困惑と下流の精度が8％以上向上します。
トークンバトラーを、新規合成の小コンテキストの共同参照検索タスクで評価し、オラクルに近い精度を実証します。
コード、モデル、およびベンチマーク：https：//github.com/abdelfattah-lab/tokenbutler

要約(オリジナル)

Large Language Models (LLMs) rely on the Key-Value (KV) Cache to store token history, enabling efficient decoding of tokens. As the KV-Cache grows, it becomes a major memory and computation bottleneck, however, there is an opportunity to alleviate this bottleneck, especially because prior research has shown that only a small subset of tokens contribute meaningfully to each decoding step. A key challenge in finding these critical tokens is that they are dynamic, and heavily input query-dependent. Existing methods either risk quality by evicting tokens permanently, or retain the full KV-Cache but rely on retrieving chunks (pages) of tokens at generation, failing at dense, context-rich tasks. Additionally, many existing KV-Cache sparsity methods rely on inaccurate proxies for token importance. To address these limitations, we introduce TokenButler, a high-granularity, query-aware predictor that learns to identify these critical tokens. By training a light-weight predictor with less than 1.2% parameter overhead, TokenButler prioritizes tokens based on their contextual, predicted importance. This improves perplexity & downstream accuracy by over 8% relative to SoTA methods for estimating token importance. We evaluate TokenButler on a novel synthetic small-context co-referential retrieval task, demonstrating near-oracle accuracy. Code, models and benchmarks: https://github.com/abdelfattah-lab/TokenButler

arxiv情報

提供元, 利用サービス

arxiv.jp, Google