Researchers from Stanford and UC Berkeley have introduced TokenGate, a component designed for pretrained LLMs that implements trainable token pruning. The method addresses the quadratic growth of attention costs by using a learnable selector that allows gradients to flow back during training, bypassing the non-differentiable nature of traditional Top-K selection.
The team evaluated the method using the Qwen3-1.7B model across five benchmarks, including RepoBench-P and needle-in-a-haystack. TokenGate demonstrated performance improvements, achieving up to 11.3x faster prefill speeds. In tests on NIAH 32k, gating at layer 3 matched full-prefill accuracy while running 3.8x faster than full prefill and 2.1x faster than FastKV.
Training TokenGate follows a TRL/PEFT recipe, with each step being 1.3x slower than standard LoRA. The researchers noted that training converged faster on four out of five tasks, with a mean training time of 1.8 A100 hours compared to 2.2 hours for baselines. The implementation requires Python 3.10+ and CUDA PyTorch 2.9+.
The system utilizes a learnable per-head gain and temperature to refine token selection. It supports flexible budgets during training by averaging full-context loss and random-budget loss. Checkpoints for the tested tasks are available on Hugging Face, and an optional FlashAttention-2 key-bias kernel is provided to accelerate inference for dense Qwen3 models.
