スタンフォード大学とカリフォルニア大学バークレー校の研究者らは、2026年10月25日、大規模言語モデル（LLM）の長文コンテキスト事前埋め込み段階向けの学習可能なトークン剪枝技術「TokenGate」を公開しました。このモジュールは既存のモデルに“プラグアンドプレイ”方式で統合され、元のタスク精度を維持したまま、最初のトークン生成時間（Time-to-First-Token, TTFT）を1.4倍から11.3倍に向上させ、ピーク時のメモリ使用量を1.5倍から3.4倍削減します。
TokenGateの核心的な仕組みは、Transformerの特定の中間層（L層と呼ぶ）の後に小型の多層パーセプトロン（MLP）によるゲートモジュールを挿入し、LoRA（r=16）を用いて軽量な微調整を行うことです。このモジュールは、入力文書の各トークンに対して重要度スコアを算出し、Top-K選択によって重要なトークンのみを保持します。その後の高層ではこれらのK個のトークンのみを処理し、それらに対するKVキャッシュは生成しません。一方、質問提示部分は強制的に保持され、剪枝の対象にはなりません。
Top-K選択が微分不可能なためエンドツーエンドでの学習ができないという問題を解決するため、TokenGateはMoEルーターに類似した仕組みを導入しました。具体的には、保留されたトークンのスコアを注意力のlogitsに対する学習可能なバイアス項に変換し、さらに各ヘッドごとのゲイン（per-head gain）と学習可能な温度係数（temperature）による補正を加えることで、勾配をスコアリングモジュールへ戻せるようにしています。学習では標準的な次単語予測損失を用い、全文脈損失とランダム予算損失を混合することで、1回の学習においてゲートモジュールが多様なプルーニング予算に対応できるようにしています。
実験はQwen3-1.7BおよびQwen3.5-2Bモデルを用いて、RepoBench-P、2WikiMultihopQA、MultiNews、BioDEX、およびneedle-in-a-haystack（NIAH）の5つのベンチマークで評価を行い、コンテキスト長は最大128kトークンに達しました。その結果、NIAH-16kタスクにおいて最優秀性能を達成し（TTFTが11.3倍に高速化、メモリ使用量が3.42倍に削減）、128kのコンテキスト条件下ではTTFTが13.0倍に高速化、メモリ使用量は約3.43倍に安定して削減されました。また、2WikiMultihopQAタスクでは、全コンテキストでの精度95%を犠牲にしながら、TTFTを3.0倍に高速化し、メモリ使用量を2.7倍に削減しました。
TokenGateは、Qwen3.5-2Bのような混合注意モデル（Gated DeltaNetと全注意層を含む）においても堅牢な性能を示します。NIAH-32kタスクでは、第3層でゲート制御を施すだけで全プリフィル精度に到達し、その速度はフルプリフィル比で3.8倍、FastKV比で2.1倍となります。さらに、第1層または第2層でゲート制御を行えば、加速効果は最大7.8倍に達し、精度の低下はわずか6%に留まります。これに対し、FastKVなどのベースライン手法は純粋な注意層にしか適用できず、DeltaNetのような非標準構造には対応できません。
この技術はオープンソースで実装されており、Python 3.10以降およびCUDA対応のPyTorch 2.9以上をサポートしています。さらに高速化のために、FlashAttention-2のキー偏置カーネルをオプションで統合することも可能です。GitHubリポジトリでは、全5タスクに対応したQwen3-1.7Bのチェックポイントに加え、NIAHタスク向けの専用ゲート制御モデル（例：Hugging FaceのモデルID：liana313/tokengate-niah-qwen3-1.7b）も提供されています。学習時間は標準的なLoRAファインチューニングの約1.3倍ですが、5つのタスクのうち4つでは収束がより早く、平均的な学習時間はA100で1.8時間（ベースラインの2.2時間と比較）です。