TokenGate speeds up large language model input processing with token pruning
TokenGate: Adapting LLMs for Fast and Memory-Efficient Prefill with Learned Token Pruning
Computation and LanguageArtificial Intelligence
Summary
Large language models read every part of long inputs, which takes a lot of time and memory. The authors introduce TokenGate, a small add-on that helps these models skip less important parts of the input, processing only what matters. This makes the models faster and use less memory without losing accuracy. TokenGate works by learning which tokens to keep and adjusts the model's attention accordingly. It can be easily added to existing language models to improve their efficiency.
What this means in practice
- •For software engineers: Optimize processing of long codebases or technical documents by reducing transformer input tokens without losing output quality.
- •For cloud ml platform operators: Reduce computational costs and memory use during large language model inference for long-context inputs by integrating TokenGate for token pruning.
Authors
Liana Patel, Rishabh Ranjan, Ion Stoica, Matei Zaharia, Carlos Guestrin
Abstract
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\times$ faster time-to-first token (TTFT) and $1.5$-$3.4\times$ 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\times$ faster prefill and $2\times$ 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: https://github.com/tokengate-models/TokenGate.