TokenGate makes AI read long documents up to 11x faster by skimming the unimportant parts

• Why it matters: When you ask an AI to read a 128,000-word document, it currently processes every single word, which is slow and expensive. TokenGate lets it skim past the irrelevant parts first, so you get answers faster and at lower cost. The code is open source for anyone to try.

What you could do with it

• A researcher could drop a small gate module into an existing LLM like Qwen3 and teach it to skim irrelevant tokens in long documents, cutting prefill time dramatically while keeping accuracy intact ¹.
• A company running AI that must digest 128,000-word inputs could slash both wait times and memory costs without retraining the entire model from scratch ².
• Someone analyzing massive code repositories could have their language model read through walls of boilerplate quickly, focusing processing power only on the sections that matter for answering their question ².

How it works

Large language models normally push every single token — roughly every word or piece of a word — through every layer of the model. The amount of work grows quadratically with input length, so long documents quickly become a bottleneck ².

TokenGate changes this by learning which tokens to keep. A small neural-gating module, combined with LoRA (a lightweight adaptation technique that adjusts only a small number of parameters), scores each token at a chosen layer. The top-scoring tokens are kept; the rest are pruned before the remaining layers process them ¹. Only the document part of the prompt is skipped — the question you ask the model is never pruned .

The trick that makes this trainable is a technique called token-gated attention. Normally, selecting the top-K tokens is not something a model can learn from because it is not differentiable — gradients cannot flow back through it. TokenGate borrows a trick from mixture-of-experts routing: each kept token's score becomes a bias on the attention calculations above the gate, giving the scorer a path to receive training signals ². The model is trained with a mix of full-context and random-budget losses, so one trained gate works across different pruning budgets ¹.

Speed and memory gains

On the Qwen3-1.7B model, the researchers measured time-to-first-token (TTFT) improvements of 1.4× to 11.3× faster than full prefill, with peak memory dropping by 1.5× to 3.4×, with no loss in accuracy ¹. On the needle-in-a-haystack benchmark at 16,000 tokens, they achieved the best result: 11.3× faster TTFT and 3.42× lower memory .

When context length grows, speedups increase further. A gate trained at 16,000 tokens reaches 12.8× faster at 32,000 tokens, 13.0× at 64,000 and 128,000 tokens, while memory savings stay near 3.43× ¹.

The team also tested Qwen3.5-2B, a hybrid-attention model that mixes standard attention layers with a different architecture called Gated DeltaNet. TokenGate works natively on both types. At layer 3, it matched full-prefill accuracy and ran 3.8× faster than full prefill and 2.1× faster than the baseline method FastKV ¹. Gating at shallower layers 1 or 2 reached 7.8× speedup while staying within 6% of full-context accuracy — an option FastKV does not offer because it can only prune standard attention layers .

Across five benchmark tasks — RepoBench-P, 2WikiMultihopQA, MultiNews, BioDEX and needle-in-a-haystack — TokenGate was Pareto-optimal against six baselines including FastKV, SlimInfer, DASH, GemFilter, FlexPrefill and MInference ¹. Against the next-best baseline, it delivered 1.2× to 6.4× faster TTFT and up to 2.1× lower memory .

How to use it

nTokenGate is released under the Apache 2.0 licence with code and checkpoints on GitHub ¹. It requires Python 3.10+ and PyTorch 2.9+ with CUDA . An optional FlashAttention-2 kernel is available for dense Qwen3 models; using it adds just 1.0× to 1.1× overhead over a plain call . The training recipe is similar to standard parameter-efficient fine-tuning, with each training step about 1.3× slower than plain LoRA; the team reported convergence in 1.8 A100 hours on average versus 2.2 for standard LoRA .

This story reports on a research release for information only. It is not advice or encouragement to use any specific tool or technique.

How good is it

TokenGate achieved 1.4× to 11.3× faster time-to-first-token and 1.5× to 3.4× lower peak memory on the Qwen3-1.7B model across five benchmarks, with no accuracy loss ¹.,On the needle-in-a-haystack benchmark at 16,000 tokens, the model ran 11.3× faster with 3.42× lower memory compared to full prefill .,When context length grew to 128,000 tokens, the same gate trained at 16,000 tokens still delivered about 13× faster prefill while memory savings stayed near 3.43× .,Compared to the best prior pruning methods, TokenGate was up to 6.4× faster and used 2× less peak memory ².

Corrections and updates

• Corrected on 27 October 2026. An earlier version said the work grows exponentially with input length; it is quadratic.

Key facts

• TokenGate adds learned token pruning to existing LLMs via a small MLP gate and LoRA (r=16)
• Prefill runs 1.4× to 11.3× faster with 1.5× to 3.4× lower peak memory, with no accuracy loss
• At 128k context a gate trained at 16k reaches 13.0× faster TTFT and ~3.43× lower memory
• Works on hybrid-attention Qwen3.5, including non-attention Gated DeltaNet layers where other methods cannot prune
• Code and checkpoints are open-source on GitHub under Apache-2.0

Sources

1. Researchers from Stanford and UC Berkeley Introduce TokenGate: Token-Gated Attention Makes Long-Context Prefill Up to 11 — marktechpost.com
2. TokenGate: Adapting LLMs for Fast and Memory-Efficient Prefill with Learned Token Pruning