Researchers from Stanford and UC Berkeley Introduce TokenGate: Token-Gated Attention Makes Long-Context Prefill Up to 11.3x Faster

This is the available feed excerpt, which may be shortened or in the publisher’s original language.

Stanford and UC Berkeley researchers released TokenGate, a module for long-context prefill. TokenGate learned token pruning lets an existing LLM skim its input. Only the tokens that matter get full processing. The research team reports 1.4x to 11.3x faster time-to-first-token (TTFT) with no accuracy loss. Peak prefill memory drops 1.5x to 3.4x. Code and checkpoints are on GitHub under Apache-2.0. TL;DR MLP gate: A small MLP gate and LoRA (r=16) added to Qwen3-1.7B or Qwen3.5-2B. Tested on contexts up to 128k tokens. Runs on: Paper timings use 1 A100 at batch size 1. Training used 4 A100s. Performance: Pareto-optimal on 5 tasks against 6 baselines. Best: NIAH 16k at 11.3x faster TTFT and 3.42x lower memory. TTFT gain hits 13.0x at 128k. Worst: Speedup drops to 1.4x on some tasks. What is TokenGate? TokenGate is a drop-in component that adds trainable token pruning to pretrained LLMs. Transformers push every token through every layer. Attention cost grows quadratically with length. Training-free pruning is cheap but loses accuracy at high pruning rates. Architectural changes often need pre-training from scratch. How does TokenGate learned token pruning work? The forward pass splits into 3 phases: Standard forward: layers below gate layer L process all N tokens as usual. Gate module: an MLP scores each token from its post-attention residual at layer L. A top-K selector keeps the highest-scoring tokens. Pruned forward: layers above L process only those K tokens and never create KV entries for pruned ones. Only the document part of the prompt is prunable. The question is never pruned. What is token-gated attention? Top-K selection is not differentiable, which blocks end-to-end training. TokenGate borrows the Mixture-of-Experts router trick. Each kept token’s score becomes a bias on attention logits above the gate. That bias gives gradients a path back to the scorer. A learnable per-head gain and a learnable temperature refine this. The temperature cor