斯坦福大学和加州大学伯克利分校的研究人员发布了 TokenGate,一个用于长上下文预填充的模块。TokenGate的可学习令牌剪枝让现有LLM能够浏览其输入,只有重要的令牌才会得到完整处理。研究团队报告了 1.4倍至11.3倍更快的首令牌时间(TTFT) 且没有准确率损失。峰值预填充内存降低1.5倍至3.4倍。代码和检查点已发布在 GitHub 上,采用Apache-2.0许可。
TL;DR
- MLP门控: 在Qwen3-1.7B或Qwen3.5-2B上添加一个小型MLP门控和LoRA(r=16)。测试的上下文长度最高达128k令牌。
- 运行环境: 论文中的计时使用1块A100、批量大小为1。训练使用了4块A100。
- 性能: 在5个任务上对比6个基线方法处于帕累托最优。
- 最佳: NIAH 16k上TTFT提速11.3倍,内存降低3.42倍。在128k时TTFT增益达到13.0倍。
- 最差: 在某些任务上加速比降至1.4倍。
什么是TokenGate?
TokenGate是一个可直接插入的组件,为预训练LLM添加可训练的令牌剪枝。Transformer让每个令牌通过每一层,注意力成本随长度呈二次增长。无需训练的剪枝虽然便宜,但在高剪枝率下会损失准确率。架构改动往往需要从头重新预训练。
TokenGate的可学习令牌剪枝是如何工作的?
前向传播分为3个阶段:
- 标准前向: 门控层L以下的各层照常处理全部N个令牌。
- 门控模块: 一个MLP根据每个令牌在L层的注意力后残差对其进行评分。top-K选择器保留得分最高的令牌。
- 剪枝前向: L层以上的各层只处理这K个令牌,并且从不为被剪枝的令牌创建KV条目。
只有提示词的文档部分可以被剪枝。问题部分永远不会被剪枝。
什么是令牌门控注意力（token-gated attention）？
Top-K 选择不可微，这阻碍了端到端训练。TokenGate 借用了混合专家（Mixture-of-Experts）路由器的技巧。每个保留令牌的分数成为注意力 logits 上高于门限的一个偏置项。该偏置为梯度提供了一条回传到打分器的路径。
一个可学习的每头增益和一个可学习的温度对此进行细化。温度用于修正每个查询看到的更少的键。训练使用普通的下一词元损失。每一步将全上下文损失和随机预算损失取平均。因此 1 个训练好的门控可以在不同预算下工作。
TokenGate 在基准测试中表现如何？
研究团队在 RepoBench-P、2WikiMultihopQA、MultiNews、BioDEX 和大海捞针（needle-in-a-haystack，NIAH）上测试了 Qwen3-1.7B。基线为 FastKV、SlimInfer、DASH、GemFilter、FlexPrefill 和 MInference，均使用相同的 LoRA 骨干。
- 对比完整预填充： TTFT 快 1.4 倍到 11.3 倍，内存低 1.5 倍到 3.4 倍。
- 对比次优基线： TTFT 快 1.2 倍到 6.4 倍，内存最多低 2.1 倍。
- 2WikiMultihopQA： 在达到 95% 全上下文准确率的情况下，预填充快 3.0 倍，内存减少 2.7 倍。
- 更长上下文： 在 16k 训练的门控在 32k 时达到 12.8 倍，在 64k 和 128k 时达到 13.0 倍。内存节省保持在约 3.43 倍。
学习到的选择器在同层优于 FastKV 的启发式方法。即使是随机 K 选择在 TokenGate 骨干上也得分更高。研究团队将此归功于骨干的协同适应。
它适用于混合注意力模型吗？
适用。Qwen3.5-2B 混合了 Gated DeltaNet 和全注意力层。在 NIAH 32k 上，第 3 层的 TokenGate 达到了与完整预填充相同的准确率。它比完整预填充快 3.8 倍，比 FastKV 快 2.1 倍。在第 1 或第 2 层进行门控达到 7.8 倍，准确率差距在 6% 以内。FastKV 只能剪枝注意力层。
TokenGate 对比 FastKV、SlimInfer、GemFilter 和 DASH
| 特性 | TokenGate | FastKV | SlimInfer | GemFilter | DASH |
|---|---|---|---|---|---|
| Token 选择 | 学习型 MLP 门控 | 来自最后查询的注意力 | 块级注意力，多阶段 | 过滤层的最后查询注意力 | 隐藏状态稳定性 |
| 是否需要训练 | 是（门控 + LoRA r=16） | 否 | 否 | 否 | 否 |
| TokenGate 的 Qwen3-1.7B 运行中的剪枝层 | 1 | 13（默认） | 8, 17, 25 | 12 | 11 |
| 在 Gated DeltaNet 层进行剪枝 | 是 | 否（仅注意力层） | 未披露 | 未披露 | 未披露 |
| 自报的最佳预填充收益（自有环境） | 11.3倍 TTFT，3.4倍内存 | 1.82倍预填充 | 2.53倍 TTFT（LLaMA3.1-8B，RTX 4090） | 相比 SOTA 速度提升 2.4 倍，GPU 内存减少 30% | 未披露 |
| 代码 | GitHub | GitHub | GitHub | GitHub | 已发布 |
自报的加速结果使用了不同的实验设置，因此不能直接比较。剪枝层结果来自 TokenGate 的附录 C.2.
如何部署 TokenGate？
该 GitHub 仓库 需要 Python 3.10+ 和 CUDA PyTorch 2.9+。一个可选的 FlashAttention-2 key-bias 内核可以加速稠密的 Qwen3 推理，其开销是普通调用的 1.0 到 1.1 倍。如果没有该内核，推理会回退到 FlexAttention。
一个 {state} 占位符标记可剪枝的文档。 generate 接受一个 max_kept_prefill_tokens 预算。检查点涵盖 Qwen3-1.7B 上的全部 5 个任务，外加一个 Qwen3.5-2B NIAH 门控。示例： liana313/tokengate-niah-qwen3-1.7b.
训练沿用了 TRL/PEFT 的配方。每一步比普通 LoRA 慢 1.3 倍。在 5 个任务中的 4 个上收敛更快。平均训练时间为 1.8 对 2.2 A100 小时。
要点
- TokenGate 通过一个小型门控和 LoRA，为现有 LLM 添加了可学习的 token 剪枝。
- 预填充速度提升 1.4 到 11.3 倍，峰值内存降低 1.5 到 3.4 倍。
- Token 门控注意力使 top-K 选择变得可训练，风格类似 MoE 路由器。
- 在 TTFT 上，它比次优的剪枝基线最高快 6.4 倍。
- 它适用于混合架构的 Qwen3.5-2B，而 FastKV 无法剪枝 DeltaNet 层。
欢迎查看 论文, GitHub 仓库 以及 Hugging Face 论文页面。另外，欢迎在 Twitter 上关注我们，别忘了加入我们的 150k+ ML SubReddit 并订阅 我们的通讯。等等！你在用Telegram吗？ 现在你也可以在Telegram上加入我们了。
需要与我们合作推广您的GitHub仓库、Hugging Face页面、产品发布或网络研讨会等吗？ 与我们联系
这篇文章 斯坦福大学和加州大学伯克利分校的研究人员推出TokenGate：令牌门控注意力使长上下文预填充速度提升至11.3倍 首发于 MarkTechPost.