TokenGate: Adapting LLMs for Fast and Memory-Efficient Prefill with Learned Token Pruning

What it claims

TokenGate, from Stanford and UC Berkeley (Liana Patel, Rishabh Ranjan, Ion Stoica, Matei Zaharia, Carlos Guestrin; arXiv, October 2026), prunes context tokens once, partway up the network. It is trained with plain next-token-prediction loss and no auxiliary loss. The authors say total training time is similar to supervised fine-tuning, with overhead close to LoRA-style parameter-efficient tuning. One fine-tuning pass reportedly generalizes across input lengths and token budgets K, so a new budget does not require retraining. Earlier learned-pruning methods that fix the budget during training do require it.

How it works

1. Layers below a chosen layer L run a normal forward pass over all N tokens.
2. At layer L, a small MLP reads each prunable token's normalized post-attention residual and emits a scalar. Scores are RMSNorm-normalized across the sequence.
3. The MLP is initialized to output zero, so training starts from the unmodified model.
4. A top-K selector keeps the K highest-scoring prunable tokens, plus all non-prunable tokens. Non-prunable tokens are answer tokens during training, decoded tokens, or tokens the caller specifies, such as the user's question.
5. Pruned tokens stop at layer L. Layers above L never process them and write no KV cache entries for them. This is where the compute and memory savings come from.
6. In the upper layers, token-gated attention multiplies each attention weight by the gate value of its key. This gives the scoring network a gradient path around the non-differentiable top-K.

flowchart TD
  A[All N tokens, layers below L] --> G[Gate MLP scores]
  G --> T{Top-K plus non-prunable}
  T -->|kept| U[Layers above L]
  T -->|pruned| X[Stopped, no KV cache]
  G -->|gate values| U

Evidence

• Qwen3-1.7B on 2WikiMultihopQA. At 95% of full-context F1, TokenGate prefills 3.0x faster with 2.7x less peak memory. It beats the training-free pruners FastKV and SlimInfer on the accuracy versus TTFT and memory Pareto frontier.
• Qwen3.5-2B (hybrid attention). TTFT is 3.8x lower than full context with no accuracy loss. It is 2.1x lower than the next-best pruning method at fixed accuracy.
• Against prior pruning methods overall. Gains reach 6.4x faster TTFT and 2x lower peak memory relative to the next-best baseline.
• Baseline for the headline range. The 1.4-11.3x TTFT figure is measured against a fine-tuned full-context model.
• Tasks. Evaluation covers code completion, multi-hop QA, multi-document summarization (MultiNews), medical extraction and needle-in-a-haystack.
• Qualitative example. With K = 256 kept tokens from a 5K-token input, a 2WikiMultihopQA answer matches the full-context answer.

How to use it

Code is at https://github.com/tokengate-models/TokenGate. Adoption requires fine-tuning the base model with the gate inserted at a chosen layer L. The fetched text gives no install commands, weights, licence or hardware figures.

Limits

• Prefill only. TokenGate targets prefill for documents supplied at request time. Decode-side KV eviction methods are described as complementary, not replaced.
• Lossy for some tasks. In the MultiNews example at K = 256, the summary drifts from the full-context one: "former Winnipeg teacher" versus "former college instructor and sports anchor".
• Small models only. The reported models are 1.7B and 2B parameters. No larger models appear in the fetched text.
• Truncated text. The fetched text stops in Section 4.2. The training recipe, choice of L and K, full results tables and any stated limitations are not shown.