Researchers from Stanford and Berkeley propose TokenGate, using learnable pruning to speed up long-context Prefill by up to about 11x
Researchers from the Stanford AI Lab and Berkeley Sky propose a pluggable small module, TokenGate, which is inserted into the shallow layers of an existing LLM and forwards only the important top-K tokens to the upper layers. The paper reports that with no loss in accuracy, time to first token is 1.4–11.3× faster and peak memory usage is 1.5–3.4× lower; compared with the best training-free pruning method, it is up to about 6.4× faster and 2× more memory-efficient. Training uses the ordinary next-token loss, with a cost close to LoRA, and after a single fine-tuning it generalizes to different input lengths and different K.
Why read this
AI score 63
The paper presents acceleration and memory data for the pluggable gating module on long-context prefill, and explains how it works with Prompt Cache.
Topics
View original postIntroducing ✨Token-Gated Language Models✨ for up to 11x faster prefill by learning to skim context, conditionally forwarding only tokens that matter. Akin to the shift from dense models to MoEs, the next frontier of efficient LLMs is conditional computation over a new scaling axis: context tokens🧵👇
View original
Introducing ✨Token-Gated Language Models✨ for up to 11x faster prefill by learning to skim context, conditionally forwarding only tokens that matter. Akin to the shift from dense models to MoEs, the next frontier of efficient LLMs is conditional computation over a new scaling axis: context tokens🧵👇
Body · AI translation
Skipping through long contexts like a human: TokenGate uses learnable pruning to speed up Prefill by up to about 11×
Researchers from @StanfordAILab and @BerkeleySky, two universities, use a pluggable small module called "TokenGate", inserted into the shallow layers of existing LLMs, to learn to "skim" long contexts: only the important top-K tokens are passed on to the upper layers. The paper reports that with no accuracy loss, time to first token can be 1.4–11.3× faster and peak memory can drop 1.5–3.4×; compared with the best training-free pruning methods, it is up to about 6.4× faster and 2× more memory-efficient. Training cost is close to LoRA, and no pretraining from scratch is needed. https://t.co/7CHPMBktuN
Background Experts reading long documents or codebases first skim and then read closely; the standard Transformer, however, computes every token at every layer. Prefill for long contexts thus becomes a bottleneck: attention grows approximately quadratically with length, while MLP and KV cache grow linearly.
The two existing approaches each have shortcomings: · Training-free pruning (FastKV, SlimInfer, etc.): relies on attention heuristic scoring, cheap to deploy, but limited speedup and prone to accuracy loss. · Architecture changes (adaptive computation, sub-quadratic attention, etc.): better results, but often require retraining and are costly.
TokenGate aims for the middle ground: lightweight adaptation on an already-trained model, learning end-to-end which tokens to prune.
Method of the paper · Forward pass in three stages: below the gating layer, all N tokens are processed as usual; the gating module scores each token and keeps only the top-K; above the gating layer, only these K are processed, and pruned tokens no longer produce upper-layer computation or KV. · Token-gated attention: top-K itself is non-differentiable. The approach is similar to MoE: the importance scores given by the gate are used to reweight the upper-layer attention logits, thereby opening up gradients for end-to-end training. · Training: ordinary next-token loss suffices, with no auxiliary loss; the authors say total training time is close to SFT / LoRA-level adaptation of the base model. After one fine-tuning, it generalizes to different input lengths and different K. · Coordination with Prompt Cache: the full KV below the gating layer can be cached normally across queries; the KV above the gating layer serves only the current question and can be discarded when the question changes. In the experiments the gate is placed at layer 1, and the specific number of layers can be adjusted according to load.
Main results (self-reported in the paper) · Task coverage: code completion, multi-hop QA, multi-document summarization, medical extraction, Needle-in-a-Haystack, etc. · Relative to the full-context fine-tuning baseline: with no accuracy loss, TTFT 1.4–11.3×, peak prefill memory 1.5–3.4×. · Relative to the strongest training-free pruning: at the same accuracy, up to about 6.4× faster and about 2× more memory-efficient. Example: 2WikiMultihopQA + Qwen3-1.7B, at 95% of full-context accuracy, about 3.0× faster and 2.7× more memory-efficient. It can also be attached to hybrid attention models (such as Qwen3.5-2B): with no accuracy loss, about 3.8× faster TTFT. · Why it works: the selection pattern varies with the task—QA keeps answer keywords, summarization keeps stopwords at the end of phrases, etc.; the backbone and gate adapt together, making it more robust to pruning.