スタンフォード大学とカリフォルニア大学バークレー校の研究者らは、長文脈のプリフィル（prefill）向けモジュール「TokenGate」を公開した。TokenGate は学習可能なトークンプルーニングによって、既存の大規模言語モデル（LLM）が入力を「ざっと読み流す」ことを可能にする。重要なトークンだけが全層での完全な処理を受ける。研究チームによれば、精度を損なうことなく、最初のトークン生成までの時間（TTFT）が 1.4 倍から 11.3 倍高速化し、プリフィル時のピークメモリが 1.5 倍から 3.4 倍削減される。コードとチェックポイントは Apache-2.0 ライセンスで GitHub に公開されている。
TL;DR（要点）
MLP ゲート：Qwen3-1.7B または Qwen3.5-2B に小さな MLP ゲートと LoRA（r=16）を追加。最大 128k トークンの文脈でテスト済み。
実行環境：論文の計測は A100 1 基、バッチサイズ 1。学習には A100 4 基を使用。
性能：5 つのタスクで 6 つのベースラインに対してパレート最適。
最良ケース：NIAH 16k で TTFT が 11.3 倍高速、メモリが 3.42 倍削減。128k 文脈では TTFT のゲインが 13.0 倍に達する。
最悪ケース：一部のタスクでは高速化が 1.4 倍まで低下する。
TokenGate とは何か
TokenGate は、事前学習済み LLM に学習可能なトークンプルーニングを追加するドロップイン（そのまま組み込み可能な）コンポーネントである。標準的な Transformer はすべてのトークンをすべての層に通すため、アテンションのコストは系列長に対して二次関数的に増大する。学習不要のプルーニングは低コストだが、高いプルーニング率では精度が低下する。一方、アーキテクチャの変更は通常、ゼロからの事前学習を必要とする。
TokenGate の学習可能なトークンプルーニングの仕組み
順伝播は 3 つのフェーズに分かれる。
標準順伝播：ゲート層 L より下の層は、通常どおり N 個のトークンすべてを処理する。
ゲートモジュール：MLP が層 L のアテンション後の残差から各トークンをスコアリングし、top-K セレクタが最もスコアの高いトークンを保持する。
プルーニング後の順伝播：L より上の層は保持された K 個のトークンのみを処理し、プルーニングされたトークンの KV エントリは一切作成しない。
プルーニング可能なのはプロンプトのうち文書部分のみで、質問は決してプルーニングされない。
トークンゲート付きアテンションとは
Top-K 選択は微分不可能であり、エンドツーエンドの学習を妨げる。TokenGate は Mixture-of-Experts（MoE）ルーターの手法を借用する。保持された各トークンのスコアは、ゲートより上のアテンション logits へのバイアスとなり、このバイアスがスコアラーへの勾配経路を与える。学習可能なヘッドごとのゲインと学習可能な温度がこれをさらに洗練する。温度は各クエリが見るキーが少なくなることを補正する。学習には通常の次トークン予測損失を用い、各ステップで全文脈損失とランダム予算損失を平均する。これにより、1 つの学習済みゲートが複数の予算で通用する。
ベンチマークでの TokenGate の性能
研究チームは Qwen3-1.7B を RepoBench-P、2WikiMultihopQA、MultiNews、BioDEX、および needle-in-a-haystack（NIAH）でテストした。ベースラインは FastKV、SlimInfer、DASH、GemFilter、FlexPrefill、MInference で、いずれも同一の LoRA バックボーン上で比較した。
完全プリフィルと比較：TTFT が 1.4 倍から 11.3 倍高速、メモリが 1.5 倍から 3.4 倍削減。
次点ベースラインと比較：TTFT が 1.2 倍から 6.4 倍高速、メモリは最大 2.1 倍削減。
2WikiMultihopQA：完全文脈の精度の 95% を達成しつつ、プリフィルが 3.0 倍高速、メモリが 2.7 倍削減。
より長い文脈：16k で学習したゲートは 32k で 12.8 倍、64k と 128k で 13.0 倍に達する。メモリ削減は約 3.43 倍で維持される。
同じ層において、学習可能なセレクタは FastKV のヒューリスティックを上回る。ランダム K 選択でも TokenGate バックボーン上ではスコアが高い。研究チームはこれをバックボーンの共適応によるものとしている。
ハイブリッドアテンションモデルでも動作するか
はい。Qwen3.5-2B は Gated DeltaNet 層とフルアテンション層を混在させている。NIAH 32k において、層 3 の TokenGate は完全プリフィルと同等の精度を達成し、完全プリフィルより 3.8 倍高速、FastKV より 2.1 倍高速に動作した。層 1 または 2 でのゲートは 7.8 倍に達し、精度差は 6% 以内だった。FastKV はアテンション層しかプルーニングできない。
TokenGate と FastKV、SlimInfer、GemFilter、DASH の比較
機能比較：
トークン選択：TokenGate は学習可能な MLP ゲート。FastKV は最後のクエリからのアテンション。SlimInfer はブロックアテンションの多段方式。GemFilter はフィルタ層での最後のクエリのアテンション。DASH は隠れ状態の安定性に基づく。
学習の必要性：TokenGate は必要（ゲート + LoRA r=16）。他の 4 手法は不要。
TokenGate の Qwen3-1.7B 実行におけるプルーニング層：113（デフォルト）。FastKV は 8、17、25。GemFilter は 12。DASH は 11。
Gated DeltaNet 層でのプルーニング：TokenGate は可能。FastKV は不可（アテンション層のみ）。SlimInfer、GemFilter、DASH は非公開。
自己申告の最良プリフィルゲイン（各設定）：TokenGate は TTFT 11.3 倍、メモリ 3.4 倍。FastKV はプリフィル 1.82 倍。SlimInfer は TTFT 2.53 倍（LLaMA3.1-8B、RTX 4090）。GemFilter は SOTA 比 2.4 倍高速、GPU メモリ 30% 削減。DASH は非公開。
コード：TokenGate、FastKV、SlimInfer、GemFilter は GitHub に公開。DASH は発表のみ。
注意：各手法の自己申告のゲインは異なる設定に基づくため、直接比較はできない。プルーニング層の情報は TokenGate の付録 C.2 による。
TokenGate のデプロイ方法
GitHub リポジトリには Python 3.10 以上と CUDA PyTorch 2.9 以上が必要。オプションの FlashAttention-2 key-bias カーネルが高密度 Qwen3 を高速化し、コストは通常呼び出しの 1.0 倍から 1.1 倍。使用しない場合、推論は FlexAttention にフォールバックする。
{state} プレースホルダがプルーニング可能な文書部分を示す。generate は max_kept_prefill_tokens 予算を受け取る。チェックポイントは Qwen3-1.7B 上の 5 タスクすべてに加え、Qwen3.5-2B の NIAH ゲートをカバーする。例：liana313/tokengate-niah-qwen3-1.7b。
学習は TRL/PEFT のレシピを踏襲し、各ステップは通常の LoRA より 1.3 倍遅い。5 タスク中 4 つでより速く収束し、平均学習時間は A100 で 1.8 時間（通常の LoRA は 2.2 時間）だった。
主なポイント
TokenGate は小さなゲートと LoRA により既存 LLM に学習可能なトークンプルーニングを追加する。
プリフィルは 1.4 倍から 11.3 倍高速化し、ピークメモリは 1.5 倍から 3.4 倍削減される。
トークンゲート付きアテンションは MoE ルーター方式で top-K 選択を学習可能にする。
TTFT において次点のプルーニングベースラインを最大 6.4 倍上回る。
ハイブリッド構成の Qwen3.5-2B でも動作し、FastKV がプルーニングできない DeltaNet 層を扱える。
論文、GitHub リポジトリ、Hugging Face の論文ページを参照のこと。