In plain words
Long prompts take time and memory to process because a language model carries every word through all its layers. TokenGate adds a small component that scores the prompt partway through and lets later layers work only on selected words. The model and this selector can be trained together.
In five task tests on Qwen3-1.7B, TokenGate reduced the wait for the first response and memory use compared with processing the full prompt, without reducing task accuracy. The results cover those tasks and tests, not every model or setup.
The research
When a question arrives alongside a long document, the model must process that document before it can begin answering. Can it reduce this initial processing time and memory use by skipping some input tokens, yet still retain the information its task requires?
TokenGate scores input tokens partway through processing and trains the model to continue with selected tokens, reducing prefill work.
TokenGate is a trainable component for existing language models. It raises a practical question: can a model process a long prompt with less time and memory while retaining the accuracy needed for its task?
The authors trained TokenGate with Qwen3-1.7B and compared it with standard full-input processing across five task families. At matched accuracy, it reduced the wait for the first response and peak prefill memory.
Those comparisons show that selecting tokens partway through processing improved efficiency in the tested tasks without an observed accuracy loss.
The evidence covers five task families on Qwen3-1.7B and one test on a Qwen3.5-2B hybrid model. Timing used batch size one on a single A100, so these results do not establish the same performance on other models, tasks, hardware or batch sizes.
Analysis
Deep
- Original paper authors
- Liana Patel, Rishabh Ranjan, Ion Stoica, Matei Zaharia, Carlos Guestrin
- Submitted to arXiv
- License
- CC BY 4.0