トークンプルーニング（LLM推論）（トークンプルーニング）

トークンプルーニング（LLM推論） — 動的トークン削除による高速化

トークンプルーニングはLLMの推論中に、Attentionスコアや活性化の大きさを基準に重要度の低いトークンをKVキャッシュから削除する技術の総称だ。長文処理（128k〜1M token）でKVキャッシュが膨大になる問題を解決し、メモリ使用量と推論速度を改善する。

なぜトークンプルーニングが有効か

TransformerのSelf-Attentionは全トークンペアの関係を計算するが、実際には:

• 局所性: 近くのトークンとの関係が遠くより重要なことが多い
• 疎性（Sparsity）: Attentionスコアは通常、少数のトークンに集中する（ヘッドごとに上位5〜10%が99%のスコアを占める）
• 情報の冗長性: 繰り返し言及される情報は一部のキートークンに集約される

これらの観察から「重要でないトークンのKVペアを捨てても品質はほぼ維持できる」という仮説が成立する。

主要手法の比較

H2O（Heavy Hitter Oracle）

「Heavy Hitter」（累積Attentionスコアが高いトークン）を保持し、それ以外を削除する:

# 累積アテンションスコアでトークンをランキング
scores = attention_weights.sum(dim=0)  # 全ヘッド・全クエリでの累積
top_k_indices = scores.topk(k=budget).indices

• 利点: 実装シンプル、既存モデルに追加訓練なしで適用可能
• 欠点: 最近のトークン（Recent tokens）の重要性を過小評価しやすい

SnapKV

プロンプト（入力）フェーズ中に重要なKVペアを選択し、デコード中は固定のKVキャッシュを使用:

• プロンプト全体を一度処理してから重要度スコアを計算
• デコード中の動的なスコア更新コストを排除
• 長いシステムプロンプト・RAGコンテキストの圧縮に特に効果的

PyramidKV

Transformerの層（Layer）ごとに異なるKVキャッシュ保持率を設定:

• 浅い層: 局所的なパターンを担うため多くを保持（80〜100%）
• 深い層: 高レベルの意味を少数のトークンで表現 → 積極的に削除（20〜40%）

層ごとの役割の違いを活用することで、均一な削除より品質を維持しながらメモリを削減できる。

効果と限界

効果: 128kトークンの長文推論でKVキャッシュメモリを30〜50%削減、スループット1.5〜2倍向上（H2O論文）。

限界:

• 削除されたトークンの情報は完全に失われる（不可逆）
• ドメインや質問タイプにより最適な保持率が異なる
• 数学・コード等、全トークンへの正確な参照が必要なタスクでは品質低下リスク