Reproduce “TokenGate: Adapting LLMs for Fast and Memory-Efficient Prefill with Learned Token Pruning” on OpenResearch
Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.
Universal app for Apple silicon and Intel. Requires macOS 11 or later.
About this paper
Liana Patel, Rishabh Ranjan, Ion Stoica, Matei Zaharia, Carlos Guestrin
TokenGate reduces long-prompt prefill work by learning which tokens to retain after an intermediate layer, then sending only those tokens through the remaining layers; importance-weighted attention lets task-loss gradients train its discrete selector. On 2WikiMultihopQA with Qwen3-1.7B, it reaches 95% of full-context accuracy with 3.0× faster TTFT and 2.7× less peak prefill memory, in batch-one prefill tests on one A100.
Summary by alphaXiv.
Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.