TokenGate：通过学习型 Token 剪枝实现 LLMs 的快速且高内存效率预填充

专业的人类读者处理大规模输入（例如代码仓库和长文档）时，会先快速浏览，仅仔细阅读重要内容。相比之下，语言模型会在每个模型层中处理所有输入 Token，导致长上下文输入产生高昂的处理成本。我们提出了 TokenGate，这是一种小型、可直接插入的组件，能够扩展现有的已训练模型，使其支持对长上下文输入进行端到端可学习的 Token 剪枝，从而实现快速且高内存效率的预填充。为在 Token 剪枝所需的不可微 top-K 算子存在的情况下实现端到端学习，我们引入了 Token-gated attention，这是一种受 Mixture-of-Experts 启发的机制，利用 Token 重要性分数对注意力 logits 重新加权。在多种任务中，TokenGate 在 transformers 的预填充阶段实现了 1.4–11.3 倍更快的首 Token 生成时间（TTFT）和 1.5–3.4 倍更低的峰值内存占用，同时不造成准确率下降。与此前表现最佳的 Token 剪枝方法相比，TokenGate 可实现最高 6.4 倍更快的预填充速度和 2 倍更低的峰值内存占用。TokenGate 原生支持扩展标准 transformers（例如 Qwen3）和混合注意力模型（例如 Qwen3.5），所需的轻量适配简单易行，开销与标准参数高效微调相当。我们的代码已发布于：https://github.com/tokengate-models/TokenGate。