TokenGate

斯坦福大学与加州大学伯克利分校的研究人员联合发布了 TokenGate，这是一个专为大语言模型（LLM）长上下文预填充阶段设计的即插即用加速模块。该模块通过引入学习化的 Token 剪枝机制，使模型能够“略读”输入内容，仅对高价值 Token 进行完整处理。

在保持模型准确性无损的前提下，TokenGate 实现了 1.4 倍至 11.3 倍的首 Token 时间（TTFT）加速，并将峰值预填充内存降低了 1.5 倍至 3.4 倍。相关代码与模型检查点已在 GitHub 上以 Apache-2.0 协议开源。

核心技术：学习化的 Token 剪枝与门控注意力

传统的免训练剪枝方法在剪枝率较高时容易导致准确性下降，而修改模型架构通常需要从头预训练。TokenGate 巧妙地通过添加一个小型 MLP 门和 LoRA（秩 r=16）模块，在现有预训练 LLM（如 Qwen3-1.7B 或 Qwen3.5-2B）中实现了可训练的 Token 剪枝。其前向传递分为三个阶段：

1. 标准前向：在门控层（Layer L）之前的网络层照常处理所有 N 个 Token。
2. 门模块评分：MLP 门根据 Layer L 的注意力后残差为每个 Token 打分，并通过 Top-K 选择器保留得分最高的 Token。值得注意的是，系统仅对提示词中的“文档”部分进行剪枝，“问题”部分始终被完整保留。
3. 剪枝前向：Layer L 之后的网络层仅处理这 K 个被保留的 Token，且绝不為被剪枝的 Token 创建 KV 缓存条目，从而大幅降低计算与内存开销。

为了解决 Top-K 选择不可微从而阻碍端到端训练的问题，TokenGate 借鉴了混合专家（MoE）路由器的技巧。它将保留 Token 的分数作为门控注意力 logits 的偏置，从而为梯度提供了一条返回评分器的路径。结合可学习的每头增益和温度参数，该机制能够使用标准的下一 Token 预测损失进行端到端训练，使单一训练的门控模块能够适应不同的计算预算。

卓越的性能表现与基准测试

研究团队在 RepoBench-P、2WikiMultihopQA、MultiNews、BioDEX 和 Needle-in-a-Haystack (NIAH) 等 5 个任务上，对 Qwen3-1.7B 进行了全面评估，并与 FastKV、SlimInfer、DASH、GemFilter 和 FlexPrefill 等 6 个基线模型进行了对比。结果显示：

• 极致加速与显存节省：相较于完整预填充，TokenGate 实现了 1.4 倍至 11.3 倍的 TTFT 加速，以及 1.5 倍至 3.4 倍的内存降低。相较于次优基线，TTFT 依然快 1.2 倍至 6.4 倍，内存节省高达 2.1 倍。
• 长上下文扩展性：在 16k 长度上训练的门控模块，在 32k、64k 和 128k 的上下文长度下，分别实现了 12.8 倍和 13.0 倍的 TTFT 加速，且内存节省稳定保持在 3.43 倍左右。
• 混合架构兼容性：TokenGate 在包含 Gated DeltaNet 和全注意力层的混合模型（如 Qwen3.5-2B）上同样表现出色。在 NIAH 32k 测试中，它在第 3 层应用门控即可匹配完整预填充的准确性，速度比完整预填充快 3.8 倍，比 FastKV 快 2.1 倍。而 FastKV 等传统方法无法对非标准注意力层（如 DeltaNet）进行剪枝。

与现有方案的差异化优势

与 FastKV、SlimInfer、GemFilter 等依赖启发式规则或特定层过滤的免训练方法不同，TokenGate 的核心优势在于其“学习化”的 Token 选择机制。实验表明，即使是在 TokenGate 骨干网络上使用随机 K 选择，其得分也高于 FastKV 的启发式选择，这证明了模型骨干与门控模块在训练过程中实现了良好的协同适应。

部署与训练指南

TokenGate 提供了极高的工程友好度：

• 环境要求：需要 Python 3.10+ 和 CUDA PyTorch 2.9+。可选配 FlashAttention-2 键偏置内核以进一步加速密集计算（开销仅为普通调用的 1.0-1.1 倍），否则将回退至 FlexAttention。
• 推理配置：在输入中使用 {state} 占位符标记可剪枝的文档部分，并在 generate 函数中通过 max_kept_prefill_tokens 参数设定计算预算。
• 训练效率：训练过程镜像了 TRL/PEFT 的配方。虽然每一步比普通 LoRA 慢约 1.3 倍，但由于其在 5 个任务中的 4 个上收敛更快，平均训练时间仅需约 1.8 个 A100 小时（对比基线为 2.2 小时）。