Researchers from Stanford and Berkeley propose TokenGate, using learnable pruning to speed up long-context Prefill by up to about 11x
English text is being prepared.
Researchers from the Stanford AI Lab and Berkeley Sky propose a pluggable small module, TokenGate, which is inserted into the shallow layers of an existing LLM and forwards only the important top-K tokens to the upper layers. The paper reports that with no loss in accuracy, time to first token is 1.4–11.3× faster and peak memory usage is 1.5–3.4× lower; compared with the best training-free pruning method, it is up to about 6.4× faster and 2× more memory-efficient. Training uses the ordinary next-token loss, with a cost close to LoRA, and after a single fine-tuning it generalizes to different input lengths and different K.
Why read this
AI score 63
The paper presents acceleration and memory data for the pluggable gating module on long-context prefill, and explains how it works with Prompt Cache.
Topics
View original postIntroducing ✨Token-Gated Language Models✨ for up to 11x faster prefill by learning to skim context, conditionally forwarding only tokens that matter. Akin to the shift from dense models to MoEs, the next frontier of efficient LLMs is conditional computation over a new scaling axis: context tokens🧵👇
View original
Introducing ✨Token-Gated Language Models✨ for up to 11x faster prefill by learning to skim context, conditionally forwarding only tokens that matter. Akin to the shift from dense models to MoEs, the next frontier of efficient LLMs is conditional computation over a new scaling axis: context tokens🧵👇
Body
像人一样跳读长上下文：TokenGate 用可学习剪枝把 Prefill 最高提速约 11 倍
来自 @StanfordAILab @BerkeleySky 两所大学的研究者们用一个可插拔的小模块「TokenGate」，插在现有 LLM 的浅层，学会对长上下文做“跳读”：只把重要的 top-K token 继续往上层传。论文报告在无精度损失时，首 token 时间可快 1.4–11.3 倍，峰值显存降 1.5–3.4 倍；相对最好的免训练剪枝方法，最高约 6.4 倍更快、2 倍更省显存。训练成本接近 LoRA，不必从零预训练。 https://t.co/7CHPMBktuN
问题背景 专家读长文档或代码库会先扫一眼再精读；标准 Transformer 却对每个 token、每一层都算一遍。长上下文的 prefill 因此成为瓶颈：注意力随长度近似平方增长，MLP 与 KV cache 线性增长。
现有两条路各有短板： · 免训练剪枝（FastKV、SlimInfer 等）：靠注意力启发式打分，部署便宜，但加速有限、精度易掉。 · 改架构（自适应计算、子二次注意力等）：效果更好，但往往要重训，成本高。
TokenGate 想走中间：在已训好的模型上轻量适配，端到端学会剪哪些 token。
论文的方法 · 前向分三段：门控层以下照常处理全部 N 个 token；门控模块给每个 token 打分，只保留 top-K；门控层以上只处理这 K 个，被剪掉的不再产生上层计算和 KV。 · Token-gated attention：top-K 本身不可微。做法类似 MoE，用门控给出的重要性分数去重加权上层注意力 logits，从而为端到端训练打通梯度。 · 训练：用普通 next-token 损失即可，无辅助损失；作者称总训练时间接近对基座做 SFT / LoRA 级适配。一次微调后，可泛化到不同输入长度和不同 K。 · 与 Prompt Cache 的配合：门控层以下的全量 KV 可跨 query 正常缓存；门控层以上的 KV 只服务当前问题，换问题可丢。实验里门控放在第 1 层，具体层数可按负载调。
主要结果（论文自报） · 任务覆盖：代码补全、多跳 QA、多文档摘要、医疗抽取、Needle-in-a-Haystack 等。 · 相对全上下文微调基线：无精度损失时 TTFT 1.4–11.3×，峰值 prefill 显存 1.5–3.4×。 · 相对最强免训练剪枝：同精度下最高约 6.4× 更快、约 2× 更省显存。示例：2WikiMultihopQA + Qwen3-1.7B，在全上下文精度的 95% 处，约 3.0× 更快、2.7× 更省显存。 也可接到混合注意力模型（如 Qwen3.5-2B）：无精度损失时约 3.8× 更快 TTFT。 · 为什么有效：选择模式随任务变化，QA 留答案关键词，摘要留短语末尾的停用词等；骨干与门控一起适配，对剪枝更鲁棒。