TokenGate adds learnable token pruning to speed long-context prefill
Researchers posted a paper on r/LocalLLaMA describing TokenGate, a small drop-in component that extends already-trained models to enable end-to-end learnable token pruning over long-context inputs for faster and more memory-efficient prefill. To learn despite the non-differentiable top-K operator, the authors introduce token-gated attention, a Mixture-of-Experts-inspired mechanism that re-weights attention logits using token importance scores. They report 1.4–11.3× faster time-to-first-token and 1.5–3.4× lower peak memory with no accuracy degradation for prefill, and up to 6.4× faster prefill and 2× lower peak memory than the best prior token-pruning methods, all figures self-reported. The authors say TokenGate natively extends standard transformers such as Qwen3 and hybrid-attention models such as Qwen3.5 with lightweight adaptation, and they released code on GitHub.
The authors note that language models push every input token through every layer, making long-context prefill costly, and present TokenGate as a drop-in way to prune less important tokens during that stage.
| Expert human readers process large inputs (e.g., code repositories and long documents) by skimming to closely read only what matters. On the other hand, language models process all input tokens through each model layer, resulting in significant processing costs for long-context inputs. We introduce TokenGate, a small, drop-in component that extends existing trained models to enable end-to-end learnable token pruning over long-context inputs for fast and memory-efficient prefill. To enable end-to-end learning despite the non-differentiable top-K operator for token pruning, we introduce token-gated attention, a mechanism inspired by Mixture-of-Experts that re-weights attention logits using token importance scores. Across diverse tasks, TokenGate achieves 1.4-11.3× faster time-to-first token (TTFT) and 1.5-3.4× lower peak memory at no accuracy degradation for prefill in transformers. Compared to the best-performing prior methods for token-pruning, TokenGate gives up to 6.4× faster prefill and 2× lower peak memory. TokenGate natively extends both standard transformers (e.g., Qwen3) and hybrid-attention models (e.g., Qwen3.5) with simple lightweight adaptation comparable to standard parameter-efficient fine-tuning. We make our code available at: this https URL. submitted by /u/pmttyji [link] [comments] |