スタンフォード大学とUCバークレーの研究者らは TokenGateを公開しました。これは長コンテキスト prefill 用のモジュールです。TokenGate の学習型トークン剪定により、既存の LLM が入力をざっと読み流すことが可能になります。重要なトークンだけが完全な処理を受けます。研究チームは、 TTFT(time-to-first-token)が1.4倍から11.3倍高速化 され、精度は損なわれないと報告しています。prefill 中のピークメモリは1.5倍から3.4倍削減されます。コードとチェックポイントは Apache-2.0 ライセンスで GitHub Apache-2.0 ライセンスの下で公開されています。
TL;DR
- MLPゲート: Qwen3-1.7B または Qwen3.5-2B に、小さな MLP ゲートと LoRA(r=16)を追加します。最大128kトークンのコンテキストでテストされました。
- 実行環境: 論文の計測タイミングはバッチサイズ1のA100 1枚で測定。学習には4枚のA100を使用。
- 性能: 5つのタスクにおいて6つのベースラインに対してパレート最適。
- 最高性能: NIAH 16k で TTFT が11.3倍高速化、メモリは3.42倍削減。128kでは TTFT の向上が13.0倍に達します。
- 最低性能: 一部のタスクでは高速化が1.4倍に低下します。
TokenGateとは?
TokenGate は、事前学習済み LLM に学習可能なトークン剪定を追加する drop-in 型のコンポーネントです。Transformer はすべてのトークンをすべての層に通します。注意機構のコストは長さに対して二次的に増加します。学習不要の剪定は安価ですが、高い剪定率では精度が低下します。アーキテクチャの変更は、多くの場合ゼロからの事前学習を必要とします。
TokenGateの学習型トークン剪定はどのように機能するのか?
フォワードパスは3つのフェーズに分かれます:
- 標準フォワード: ゲート層Lより下の層は、通常どおりすべてのNトークンを処理します。
- ゲートモジュール: MLP が層Lの attention 後の残差から各トークンをスコアリングします。top-K セレクタが最もスコアの高いトークンを保持します。
- 剪定後のフォワード: 層Lより上の層は、そのK個のトークンのみを処理し、剪定されたトークンの KV エントリは決して作成しません。
剪定できるのはプロンプトのドキュメント部分のみです。質問は決して剪定されません。
トークンゲート・アテンションとは何か？
Top-K選択は微分不可能であり、エンドツーエンドの学習が妨げられます。TokenGateはMixture-of-Expertsのルーターのトリックを借用します。保持された各トークンのスコアは、ゲートより上のアテンションのロジットに対するバイアスになります。このバイアスが、スコアラーへの勾配の経路を提供します。
学習可能なヘッドごとのゲインと学習可能な温度がこれを洗練します。温度は、各クエリがより少数のキーを見ることを補正します。学習には通常のnext-token損失を使用します。各ステップで、フルコンテキスト損失とランダム予算損失を平均します。そのため、1つの学習済みゲートが複数の予算で機能します。
TokenGateはベンチマークでどの程度の性能を示すのか？
研究チームはQwen3-1.7BをRepoBench-P、2WikiMultihopQA、MultiNews、BioDEX、needle-in-a-haystack（NIAH）でテストしました。ベースラインは、同じLoRAバックボーン上でFastKV、SlimInfer、DASH、GemFilter、FlexPrefill、MInferenceです。
- フルprefillとの比較： TTFTが1.4倍から11.3倍速く、メモリが1.5倍から3.4倍少ない。
- 次点のベースラインとの比較： TTFTが1.2倍から6.4倍速く、メモリは最大2.1倍少ない。
- 2WikiMultihopQA： フルコンテキスト精度の95%で、prefillが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はフルprefillの精度に一致しました。フルprefillより3.8倍速く、FastKVより2.1倍速く実行されました。層1または層2でのゲーティングは7.8倍に達し、精度は6%以内の差でした。FastKVはアテンション層しかプルーニングできません。
TokenGate vs FastKV、SlimInfer、GemFilter、DASH
| 機能 | TokenGate | FastKV | SlimInfer | GemFilter | DASH |
|---|---|---|---|---|---|
| トークン選択 | 学習されたMLPゲート | 最後のクエリからの注意機構 | ブロック注意、マルチステージ | フィルタ層でのラストクエリ注意 | 隠れ状態の安定性 |
| 学習の要否 | 必要(ゲート + LoRA r=16) | 不要 | 不要 | 不要 | 不要 |
| TokenGateのQwen3-1.7Bで実行される剪定層 | 1 | 13(デフォルト) | 8, 17, 25 | 12 | 11 |
| Gated DeltaNet層で剪定 | はい | いいえ(注意層のみ) | 非公開 | 非公開 | 非公開 |
| 自己申告の最高プリフィル向上(独自環境) | TTFT 11.3倍、メモリ 3.4倍 | プリフィル 1.82倍 | TTFT 2.53倍(LLaMA3.1-8B、RTX 4090) | SOTA比で速度 2.4倍、GPUメモリ 30% 削減 | 非公開 |
| コード | GitHub | GitHub | GitHub | GitHub | 発表 |
自己申告による改善結果は異なる設定で得られたものであるため、直接比較はできない。プルーニング層は以下から引用した。 TokenGateの付録C.2.
TokenGateはどのようにデプロイするのですか?
The GitHub リポジトリ Python 3.10+ および CUDA PyTorch 2.9+ が必要です。オプションの FlashAttention-2 key-bias カーネルは、dense Qwen3 の処理を高速化します。そのコストは通常の呼び出しの 1.0 倍から 1.1 倍です。これがない場合、推論は FlexAttention にフォールバックします。
A {state} placeholderは剪張可能なドキュメントを示します。 generate 〜を取る max_kept_prefill_tokens budget. チェックポイントは、Qwen3-1.7B の全 5 タスクをカバーし、さらに Qwen3.5-2B の NIAH ゲートも含みます。例: liana313/tokengate-niah-qwen3-1.7b.
学習はTRL/PEFTのレシピを模倣しています。各ステップは通常のLoRAより1.3倍遅くなります。5タスク中4タスクでより早く収束しました。平均学習時間は2.2 A100時間に対して1.8 A100時間でした。
重要ポイント
- 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 論文ページ. また、ぜひ以下で私たちをフォローしてください Twitter そして、私たちの 150k以上のML SubReddit また、 弊社のニュースレターを購読してください。お待ちください!Telegramを使っていますか? 現在、Telegramでも参加できるようになりました。
GitHubリポジトリ、Hugging Faceページ、製品リリース、ウェビナーなどのプロモーションで弊社と提携したい方は? ぜひご連絡ください
本記事 StanfordとUC Berkeleyの研究者らがTokenGateを発表:トークンゲート式アテンションにより長コンテキストPrefillが最大11.3倍高速化 は最初に MarkTechPost.