TokenGate：用可学习 token 剪枝加速长上下
来自 arXiv · 期号 2026-10-20
摘要
大模型处理长上下文时，所有 token 都要经过每一层，计算开销很大。 TokenGate 是一个轻量插件，给现有模型加一个可学习的 token 剪枝机制：在注意力计算前，用重要性分数给 token 排序，只保留最重要的 K 个。为了绕过 top-K 操作不可导的问题，它借鉴了 MoE 的思路，用门控注意力重新加权注意力 logits。 在多种任务上，TokenGate 让首 token 生成时间（TTFT）快 1.4-11.3 倍，峰值内存降低 1.5-3.4 倍，且精度无损。相比此前最好的 token 剪枝方法，TTFT 最多快 6.4 倍，内存最多降 2 倍。它可以直接适配 Qwen3 等标准 Transformer 和 Qwen3.5 等混合注意力模型，微调成本与参数高效微调相当。 论文附了开源代码。
深读
导读 大模型处理长上下文时，所有输入 token 都要经过每一层，计算开销随序列长度线性增长。TokenGate 是一个轻量插件，通过可学习的 token 剪枝机制，在注意力计算前只保留最重要的 K 个 token，使首 token 生成时间（TTFT）快 1.4–11.3 倍、峰值内存降低 1.5–3.4 倍且精度无损。读完能理解一种绕过 top-K 不可导问题的实用方法，以及它如何直接适配 Qwen3 等现有模型。 关键信息 - TokenGate 在多种任务上实现首 token 生成时间（TTFT）加速 1.4–11.3 倍，峰值内存降低 1.5–3.4 倍，且无精度损失。 - 相比此前最好的 token 剪枝方法，TokenGate 的 TTFT 最多快 6.4 倍，峰值内存最多降低 2 倍。 - 该方法可原生扩展标准 Transformer（如 Qwen3）和混合注意力模型（如 Qwen3.5），微调成本与参数高效微调相当。 - 论文代码已开源：https://github.com/tokengate-models/TokenGate。 逐节讲解 问题：长上下文的计算瓶颈 人类阅读长文档或代码仓库时会先扫读、再精读关键部分。而语言模型对所有输入 token 逐层处理，长上下文导致巨大的计算开销。论文以此类比引出核心问题：能否让模型也学会只处理重要的 token？ 方法：可学习的 token 剪枝与门控注意力 TokenGate 是一个即插即用的小组件，扩展已有训练好的模型，实现端到端可学习的 token 剪枝。核心挑战在于：按重要性分数取 top-K 的操作不可导，无法直接做梯度回传。论文借鉴 Mixture-of-Experts（MoE）的思路，提出 token-gated attention：用 token 重要性分数重新加权注意力 logits，从而绕过不可导的离散选择，使整个剪枝过程可端到端训练。 效果与适配性 在多种任务上，TokenGate 使 TTFT 加速 1.4–11.3 倍，峰值内存降低 1.5–3.4 倍，且精度无损。相比此前最好的 token 剪枝方法，TTFT 最多快 6.4 倍，内存最多降 2 倍。该方法可直接适配标准 Transformer（如 Qwen3）和混合注意力模型（如 Qwen3.5），微调成本与参数高效微调相当。 值得注意 论文只有摘要，没有提供具体实验设置、任务类型、模型规模、剪枝比例等细节。TTFT 加速和内存降低的上下限（1.4–11.3 倍、1.5–3.4 倍）跨度较大，实际效果可能高度依赖任务和配置。此外，摘要未说明 token 重要性分数的计算方式，以及剪枝对生成质量（而非仅预填充阶段）的潜在影响。