技術概要：TokenButler

問題提起

大規模言語モデル（LLM）は、コンテキスト長が 64k〜1M トークンに拡張されるにつれ、メモリと計算の重大なボトルネックに直面しています。効率的なデコーディングのためにトークンの履歴を格納するキー・バリュー（KV）キャッシュは、コンテキスト長に比例して線形に成長し、膨大な GPU メモリと帯域幅を消費します。

このボトルネックを緩和する既存の解決策は、一般的に以下の 3 つのカテゴリに分類され、それぞれ固有の限界を有しています：

1. 静的戦略（例：StreamingLLM）： スライディングウィンドウなどの固定ルールを用いてキャッシュを制限します。これらはクエリを認識せず、動的な入力ニーズに適応できません。
2. 適応的エビクション（例：H2O、SnapKV）： アテンションスコアなどの指標に基づき、重要度が低いと判断されたトークンを永続的に削除します。これは、重要度が非単調であることが多いため、後で参照される可能性のある重要なトークン（初期に言及されたエンティティが後ほど参照される共参照コンテキスト）を失うリスクがあります。
3. 適応的動的戦略（例：Quest）： 完全なキャッシュを維持しつつ、デコード時にトークンのサブセット（ページ）を取得します。永続的なエビクションを回避しますが、粗粒度のページレベルの取得に依存しているため、ページ境界をまたいで重要な情報が分割されている場合や、微細な選択を必要とする密でコンテキストに富んだタスクでは失敗します。

さらに、多くの既存の手法は、トークンの重要度に対する不正確な代理指標（例：最近性や集約アテンションスコア）に依存するか、すべてのトークンに対して完全なアテンションスコアを計算する際に高い計算オーバーヘッドを伴います。

手法：TokenButler

著者らは、トークンを永続的にエビクションしたり、粗粒度のページングに依存したりすることなく、重要なトークンを特定するために設計された、高粒度かつクエリを認識する予測器「TokenButler」を提案します。

コアアーキテクチャ

TokenButler は、固定された深度ストライド（プロデューサー層）で低次元の重要度クエリを予測し、それらを後続の層（コンシューマー層）における実際の KV キャッシュキーの学習済み射影と組み合わせることで動作します。

• クエリ予測： 軽量なマルチレイヤーパーセプトロン（MLP）が「プロデューサー」層からの隠れ状態を受け取り、アテンションヘッドごとに $G$ 個のスロット固有のクエリベクトルを生成します。これらのベクトルは低次元（$d' \ll E$）です。
• キーキャッシュ射影： 各コンシューマー層において、実際の高次元 KV キャッシュキーは、学習済み行列を用いて同じ低次元空間に射影されます。
• スコアリング： 予測されたクエリを射影されたキーと内積することで、トークンの重要度スコアを安価に生成し、固定予算下での動的なトークン選択を可能にします。

学習

• 蒸留： TokenButler は、固定されたベース LLM のマスク化された因果アテンション分布を蒸留することで学習されます。目的は、教師（ベース LLM）のアテンション対数オッズと生徒（TokenButler）の予測対数オッズとの間の交差エントロピー損失を最小化することです。
• 効率性： 学習中に長いコンテキストを処理するために、損失はクエリ位置のサブサンプルセット上で計算され、後期コンテキスト領域と最終トークンに焦点を当てます。これにより、計算複雑性は $O(L^2)$ から $O(RL)$ に削減されます。
• オーバーヘッド： 予測器は、ベースモデルに対して 1% 未満のパラメータオーバーヘッドしか追加しません。

推論戦略

TokenButler は、メモリアクセスを最適化するために KV キャッシュを 3 つのバッファに分割します：

1. シンクバッファ： 最初の $S$ トークン（常に保持）。
2. ローカルウィンドウバッファ： 最新の $N$ トークンの循環バッファ（常に保持）。
3. 重要バッファ： TokenButler によって予測スコアに基づいて動的に埋められるバッファ。

最適化：

• バッチ処理： キーの射影は、各デコードステップではなく $N$ トークンごとのバッチで実行され、時間的局所性を活用します。
• 予測間隔と近隣取得： 予測器のコストを償却するため、システムは各ステップで予測器を実行しません。代わりに、$N$ ステップごと（予測間隔）に実行します。間隔の間に重要度が変化するトーク見逃すリスクを軽減するため、システムは「近隣取得」を採用し、選択されたトークンに空間的近隣を含めるように拡張します。これにより、複数トークンのエンティティや推論チェーンの網羅が保証されます。

主な貢献

1. 軽量予測器： アテンション蒸留を通じて学習され、最小限のパラメータオーバーヘッド（<1%）で微細な推定を提供する、トークンの重要度に関する学習済み指標。
2. 合成ベンチマーク： 既存のエビクションおよびページベースの手法が、永続的なトークン損失またはページミスにより、短いコンテキスト長（300 トークン）であっても失敗することを示す、新しい共参照検索タスク。TokenButler はこのタスクにおいて、ほぼオラクルレベルの精度を達成します。
3. 長コンテキスト性能： RULER および LongBench ベンチマーク（64K コンテキスト）において、最先端のスパース化および圧縮手法と比較して、競合するまたは優れた結果を示します。
4. 効率性の向上：
   • 予測間隔を使用することで、Dense Attention に対して最大 1.6 倍のオン GPU スピードアップ を達成。
   • CPU オーフローディングを伴う Dense Attention と比較して最大 7.6 倍のレイテンシ削減 を達成し、メモリ制限のあるシステムでの効率的なデコーディングを可能に。
   • 予測器の呼び出し回数を 16 倍削減しても、精度をステップごとのベースラインの 約 1.1% 以内 に維持。

実験結果

• 合成共参照タスク： TokenButler は Llama-3.1-8B および Qwen2.5-7B において約 80-95% の精度とカバレッジを達成し、50% 未満の精度に低下したエビクションベースの手法（H2O、SnapKV）や、部分的なページミスに苦しんだページベースの手法（Quest）を大幅に上回りました。
• RULER ベンチマーク（64K コンテキスト）： TokenButler は Llama-3.1-8B および Qwen2.5-7B において、SnapKV、xKV、Quest などのベースラインと同等またはそれ以上の性能を示し、Qwen2.5-7B において最良の平均スコアを達成しました。
• LongBench： TokenButler は、さまざまなデータセットにおいて、従来の KV スパース化および圧縮手法よりも高い精度を示しました。
• 推論モデル（AIME24）： DeepSeek-R1-Distill-Llama-8B において、TokenButler は 30.0% の精度を達成し、Dense デコーディングベースライン（33.3%）に迫り、同じ 8K トークン予算下で他のスパース化手法を上回りました。
• レイテンシ： CPU オーフローディングシナリオ（256K 以上のコンテキスト）において、TokenButler はトークンあたりのレイテンシを Dense 方式の約 3.2 秒から TokenButler 方式の約 0.6 秒に削減しました。

意義

本論文は、トークンの重要度は極めて動的であり、クエリに依存するため、静的ヒューリスティックや粗粒度の検索は、複雑な長コンテキストタスクには不十分であると主張しています。TokenButler は、学習された微細な選択メカニズムがこれらのヒューリスティックを効果的に置き換え得ることを実証しています。完全な KV キャッシュを維持しつつ、最も関連性の高いトークンのみを選択的にアクセスすることで、TokenButler は永続的エビクションに伴う品質低下を招くことなく、メモリ帯域幅のボトルネックを緩和します。提案された予測間隔および近隣取得戦略は、そのような学習済み予測器の計算コストを効果的に償却可能であることを示しており、限られたハードウェア上でのリアルタイム長コンテキスト推論に対する実用的な解決策となります。