TokenGate：可学习词元剪枝让大模型高效预填充

TokenGate: Adapting LLMs for Fast and Memory-Efficient Prefill with Learned Token Pruning
Liana Patel、Rishabh Ranjan、Ion Stoica、Matei Zaharia、Carlos Guestrin（Stanford University、UC Berkeley）
arXiv:2610.26754v1 [cs.CL]，提交于 2026-10-19 · 论文 · 可访问代码仓库 token-gate-llms/TokenGate（Apache-2.0）
代码状态：论文印出的 tokengate-models/TokenGate 链接与实际公开仓库不一致；本轮按 README、训练配置和关键源码核对的是 token-gate-llms/TokenGate，默认分支为 main，最近一次核对的提交是 6e51aac5。

一句话总结

TokenGate 在 Transformer 的某一层插入一个小型评分门，保留最高分的 $K$ 个文档词元和非剪枝词元，只让它们继续经过上层；它用“词元门控注意力”把不可导的 top-$K$ 选择变成可端到端训练的任务损失路径，在 Qwen3-1.7B 和 Qwen3.5-2B 上把长上下文预填充的首词延迟降低 $1.4$–$11.3\times$，并在精度不降时把峰值显存压低 $1.5$–$3.4\times$。

一、研究背景：预填充是长上下文服务的共同瓶颈

大语言模型处理代码仓库、长文档或检索结果时，用户真正等待的第一段时间并不是逐词解码，而是模型先把全部输入扫过各层并建立 KV 缓存。这个阶段通常称为预填充。论文给出的成本结构很直接：注意力计算随输入长度近似二次增长，MLP 计算与 KV 缓存随长度线性增长。即使模型在生成阶段做了 KV 驱逐或稀疏解码，只要预填充仍把所有文档词元完整送过每一层，首词延迟和瞬时显存就会继续限制服务容量。

论文从人类阅读行为得到类比：面对很长的材料，读者会略读，再把注意力集中到可能影响答案的片段。这个类比不能直接证明词元可以安全删除，却提出了一个明确的工程目标：让模型自己学习哪些输入值得进入较深层。若在网络的某一中间层完成选择，下层仍可提供上下文混合后的表示，上层则只处理小规模词元集合，计算和缓存都能同时下降。

现有免训练词元剪枝方法使用注意力分数等启发式信号，部署成本低，但在激进预算下容易出现精度断崖。其困难不只是“分数是否准确”，还包括被保留词元与模型主干没有共同适配：底层模型仍然假定全部词元存在，到了中层突然少掉一部分，后续层只能在新的统计分布上被动工作。TokenGate 的中心判断是，选择和适配必须一起学习；只改选择器而不训练主干，会留下无法由启发式修复的分布偏移。

另一条路线是重新设计 Transformer 或预训练稀疏结构。它们可能获得更彻底的效率收益，却往往需要从零预训练、特殊硬件支持或大规模架构改动。TokenGate 选择更保守的增量方案：保留原模型主体，只在指定层增加评分分支，并以 LoRA 和一个轻量门控模块完成适配。论文报告的平均训练总时长甚至略低于普通 SFT，说明它希望进入的是现有模型的微调流程，而不是另起一种基础模型架构。

这篇论文没有机器人实验。把它放进 RobotWorld 的语境时，合理的作用是“潜在基础设施”：视觉语言动作模型、长视频理解和多传感器日志都可能产生数千到数十万词元，若其中大量背景词元对当前指令无用，学习式预填充剪枝可能降低端侧或服务端等待时间。不过这一迁移目前只是待验证假设，不能把语言模型上的 TTFT 数字直接当作机器人闭环控制收益。

二、问题形式化：请求指令与请求时文档

论文把输入写成 $x=[x_R,x_D]$。其中 $x_R=(x_1,\ldots,x_r)$ 是问题、指令和回答模板，$x_D=(x_{r+1},\ldots,x_N)$ 是请求时才到达的文档。这个划分很重要：文档没有可预先计算的 KV 缓存，所以核心成本发生在一次预填充中；同时，问题本身通常很短，却决定哪些文档内容有意义，因此不能和文档一起被任意剪掉。

模型先把每个词元嵌入为 $h_i^0$，再逐层更新。标准 Transformer 层先做注意力和残差，再做 MLP 与残差，可写为

$$a_i^{\ell}=h_i^{\ell}+\operatorname{Attn}^{\ell}\!\left(\operatorname{LN}_1(h^{\ell})\right)_i,\qquad h_i^{\ell+1}=a_i^{\ell}+\operatorname{MLP}^{\ell}\!\left(\tilde a_i^{\ell}\right),\qquad \tilde a_i^{\ell}=\operatorname{LN}_2(a_i^{\ell}) \tag{1}$$

论文把 $a_i^\ell$ 称为第 $\ell$ 层词元 $i$ 的“注意力后残差”。TokenGate 的评分器正是读取门层上的这个表示，因为它已经混合了该词元此前能看到的信息，却还没有经过本层 MLP。与直接使用浅层词嵌入相比，这个位置给评分器更多上下文；与放在更深层相比，它又为上层节省了更多计算。

标准因果注意力为

$$\operatorname{Attn}_i=\sum_{j\le i}\alpha_{ij}v_j,\qquad \alpha_{ij}=\operatorname{softmax}_{j\le i}(z_{ij}),\qquad z_{ij}=\frac{q_i^{\top}k_j}{\sqrt{d_{\mathrm{head}}}} \tag{2}$$

在完整预填充中，$N$ 个词元都要通过每层，并在每层缓存键和值。TokenGate 改变的是从这个中间层开始可见的词元集合，不只是减少注意力权重中的一个小项：被剪掉的表示不再进入上层 MLP，也不产生上层 KV 条目。因此它同时触及注意力的 $O(N^2N_L)$ 项、MLP 的 $O(NN_L)$ 项和缓存的 $O(NN_L)$ 项，但门层以下的完整前向仍然存在，这一结构性事实也解释了后文显存收益为何有上限。

三、方法：评分、选择与门控注意力

3.1 评分网络与不可导的 top-K

令 $\mathcal P$ 为可剪枝的文档词元集合，$\mathcal N$ 为不可剪枝集合，包含问题、指令、回答模板，训练时还可包含监督目标。一个 MLP 对每个 $j\in\mathcal P$ 的归一化表示打分，再沿序列轴做 RMSNorm，得到可跨请求比较的重要度 $b_j$：

$$\tilde b_j=\operatorname{MLP}_{\theta}(\tilde a_j^L),\qquad b_j=\frac{\tilde b_j}{\sqrt{\frac{1}{|\mathcal P|}\sum_{j'\in\mathcal P}\tilde b_{j'}^2+\epsilon}},\qquad b_j:=0\quad(j\in\mathcal N) \tag{3}$$

这里对序列而非特征维做归一化，是为了让一次请求内部的相对重要度稳定。MLP 初始化为零输出，因此训练起点时所有文档分数相同，模型从“无偏好”开始学习。随后选择集合

$$S=\operatorname{topK}_{j\in\mathcal P}(b_j)\cup\mathcal N \tag{4}$$

被保留的 $S$ 才进入门层 MLP 与更高层。若只看这条路径，梯度会在 top-$K$ 离散选择处断开。论文没有使用近似梯度直接穿过选择器，而是让分数在注意力内部继续产生影响，这使模型仍可通过任务损失学习“哪些词元值得被选择”。

3.2 从 MoE 路由借来的可训练门控

最基础的词元门控注意力在注意力权重上乘以每个键的门值 $g_j$：

$$\alpha_{ij}=\frac{\exp(z_{ij})g_j}{\sum_{j'\in S,j'\le i}\exp(z_{ij'})g_{j'}},\qquad g_j=\operatorname{softmax}_{j\in S}(b_j) \tag{5}$$

这与混合专家模型有形式上的相似：MoE 用路由权重缩放被激活专家的输出，TokenGate 则用门值缩放被保留词元对注意力的贡献。区别在于稀疏轴从“参数”换成“输入位置”。只要预算 $K>1$，任务损失相对于被保留词元的 $b_j$ 通常有非零导数；能够降低损失的关键词元获得更高权重，干扰项获得更低权重。

直接把同一组 $b_j$ 用作所有层的门值过于僵硬。论文先为每层、每个注意力头学习正增益 $\gamma^{\ell h}$。

$$g_j^{\ell h}=\operatorname{softmax}_{j\in S}\!\left(\operatorname{softplus}(\gamma^{\ell h})b_j\right) \tag{6}$$

softplus 保证增益为正，避免高分词元因为学出负增益反而被压到低权重。初始化时，有效增益与对应层注意力 logit 的标准差对齐，使新增偏置与原模型数值尺度相近。头级参数也让不同注意力头选择不同敏感度：某些头可能偏向实体，另一些头可能偏向代码符号。

第二个修正是温度重标定。剪枝后，一个查询可见的键数量从原来的位置 $p_i$ 降为保留集合中的名次 $m_i$，softmax 往往比全上下文更尖锐。论文构造

$$\tau_i^{\ell h}=\operatorname{softplus}(\sigma^{\ell h})\frac{\log m_i}{\log\max(2,p_i)} \tag{7}$$

这个对数比值随保留比例变化，训练后能够把不同预算下的注意力温度拉回较稳定的范围。综合头级偏置与温度后：

$$\alpha_{ij}^{\ell h}=\operatorname{softmax}_{j\in S,j\le i}\left[\tau_i^{\ell h}\left(\frac{q_i^{\top}k_j}{\sqrt{d_{\mathrm{head}}}}+\operatorname{softplus}(\gamma^{\ell h})b_j\right)\right] \tag{8}$$

式 (8) 同时保留因果掩码、词元重要度偏置和查询相关温度。它可以用每键偏置的 Fused FlashAttention 变体实现；附录报告，带偏置调用相对于普通调用约多 $1.0$–$1.1\times$ 的成本，而不是退回完全未融合的软件注意力。

3.3 训练目标：全预算锚点加随机预算

训练若只采一个很小的 $K$，模型可能牺牲全上下文能力；若只训练全上下文，评分器又得不到剪枝压力。论文因此采两个预算。全预算使用全部可剪枝词元，剪枝预算从对数均匀分布抽取：

$$K_{\mathrm{full}}=|\mathcal P|,\qquad K_{\mathrm{prune}}\sim\operatorname{LogUniform}\left(1,|\mathcal P|/2\right),\qquad \mathcal L=\frac{1}{2}\left(\mathcal L_{K_{\mathrm{full}}}+\mathcal L_{K_{\mathrm{prune}}}\right) \tag{9}$$

单个预算的任务损失仍是标准下一词交叉熵：

$$\mathcal L_K=-\sum_{t\in\mathrm{target}}\log p_{\phi}\left(x_t\mid x_{<t};K\right) \tag{10}$$

可训练参数 $\phi$ 包括评分网络、增益、温度、门控分支中的 LoRA，以及主干上的标准 LoRA 参数。对数均匀采样覆盖从极小预算到一半上下文的尺度，使一个 checkpoint 在推理时可改变 $K$，无需为每个预算重新训练。工程实现把两个目标打包为变长前向，门层以下只跑一次，两个目标共享计算图。

flowchart TD
    A[请求指令 x_R + 文档 x_D] --> B[第 0 至 L-1 层
完整预填充全部 N 个词元]
    B --> C[第 L 层注意力后残差 a_i^L]
    C --> D[克隆门层 MLP 的评分分支
序列 RMSNorm 得到 b_j]
    C --> E[不可剪枝集合 N
问题、指令、回答模板]
    D --> F[可剪枝文档集合 P 上取 top-K]
    F --> G[保留集合 S = topK(P) ∪ N]
    E --> G
    G --> H[截断隐藏状态并重建位置索引]
    H --> I[第 L 层 MLP 与第 L+1 以上层
每层每头加入温度 τ 和偏置 g]
    I --> J[只缓存保留词元的上层 KV]
    J --> K[生成首词及后续解码]
    D -. 重要度经门控 logits 反传 .-> I

3.4 代码实现与论文公式的对应关系

公开仓库的最终门模块固定了论文需要的关键结构：序列轴 RMSNorm、softplus 正增益、可学习温度、共享读出、每层每头偏置，以及从门层 MLP 克隆的评分分支。tokengate/e2e_gate/gate_module/e2e_gate_module.py:204-239 把这些配置写在 E2EGate 中；同文件 :278-305 建立原始增益与温度参数，并在前向时用 softplus 转为正数。

实际选择路径在 tokengate/e2e_gate/gate_module/gate_forward.py:1155-1207。下面摘出与式 (3)、(4) 最直接对应的几行：

b = gate(h_art_cast)
K_clamped = max(1, min(K, n_art))
b_scores = gate.aggregate(b)
topk_local = _select_indices(b_scores, K_clamped, gate)
keep_indices = _keep_indices(topk_local, article_start, article_end, seq_len, hs_device)
hs_pruned = hidden_states[:, keep_indices, :]

这段代码不是示意伪代码，而是把“只对文章词元评分、保留固定集合、再按全序列索引剪裁隐藏状态”的真实顺序写清楚。随后的 :1200-1208 根据保留位置重算旋转位置编码，并把门控偏置安装到 Phase 2。对 Qwen3.5 的 FlexAttention 路径，tokengate/e2e_gate/attn/qwen3_flex_attention.py:584-628 把学习温度折入查询缩放，把每键偏置传入带门控的注意力 kernel，对应式 (8) 的实现方式。

基准脚本 benchmarks/run_benchmark.py:1-12 明确区分 full 基线、TokenGate 全预算和 TokenGate@K 三个分支；:117-123 用批大小为 1 逐个样本计时。这个设计对单用户时延更忠实，但不能代表多租户批处理吞吐，读结果时必须把这一实验边界一起带上。

四、实验结果：更快预填充与更小峰值显存

论文在五个任务上评估：RepoBench-P 的下一行代码补全、2WikiMultihopQA 的多跳问答、MultiNews 的多文档摘要、BioDEX 的生物医学药物反应抽取，以及 NIAH-16k 的针在草堆检索。主干包括标准注意力 Qwen3-1.7B 与每四层含一次全注意力的混合架构 Qwen3.5-2B；对比对象包括 FastKV、SlimInfer、DASH、GemFilter、FlexPrefill 与 MInference。

主结果的总体区间是：相对完整上下文，TTFT 降低 $1.4$–$11.3\times$，峰值预填充显存降低 $1.5$–$3.4\times$，且报告为精度不降；相对同精度下最好的免训练基线，论文声称最多 $6.4\times$ 更快和约 $2\times$ 更省显存。2WikiMultihopQA 的代表点更具体：当保留全上下文精度的 $95\%$ 时，TokenGate 比标准预填充快 $3.0\times$，峰值显存少 $2.7\times$。

图 3 的横轴只有方法真正跑出的时间，而不是按未处理词元数反推的理论 FLOPs。尤其值得注意的是它同时画 FastKV 的默认层与第 1 层。第 1 层与 TokenGate 门层相近，避免把“TokenGate 选了一个更有利的深层位置”误当成训练机制本身的收益。论文观察到即使在同层比较，可学习选择器仍明显领先，这支持其“端到端协同适配”而非单纯“更多训练算力”的解释。

图 5 进一步把收益拆成两个来源。一方面，同样使用 TokenGate 主干，学习到的 top-$K$ 在所有预算上优于随机 top-$K$；另一方面，同样使用随机选择，TokenGate 主干又优于只做 LoRA 的原始主干。前者说明评分网络确实提供了任务相关信息，后者说明主干在被训练时改变了表示分配，对偶然保留的词元更鲁棒。两个效应叠加后，才形成相对于 FastKV 选择器的优势。

图 2 展示了真实样本上的词元选择。2WikiMultihopQA 中，门控更集中保留实体、年份和答案相关短语；MultiNews 中则出现另一类模式，保留较多功能词与分散事件词。作者据此强调选择是任务依赖的，而不是统一的“专有名词抽取器”。这也是固定启发式难以跨代码、问答、摘要和医学抽取同时保持 Pareto 最优的原因。

4.1 混合注意力架构

Qwen3.5-2B 每四层使用一次全注意力，其余层为 GatedDeltaNet，没有标准注意力权重可被 FastKV 的注意力启发式直接利用。TokenGate 的评分来自门层残差，门控偏置可通过 FlexAttention 注入全注意力层，因此可以放在两类层上。论文在 NIAH-32k 上比较不同层位：门放在第 3 层时，相对标准预填充有 $3.8\times$ 更低 TTFT 且精度不降；在相同精度下比最好的 FastKV 运行快 $2.1\times$。把门前移到第 1、2 层可把 TTFT 降幅提高到 $7.8\times$，代价是相对标准预填充约 $6\%$ 的精度损失。

这张图的解释不能简化成“层越浅越好”。浅层门越早截断词元，计算节省越大，却更可能丢掉尚未来得及与问题语义对齐的实体；第 3 层精度持平，第 1、2 层时间更短但有精度预算。对部署而言，门层是速度与质量的显式旋钮；论文的默认选择并不是盲目追求最大倍数。

4.2 更长上下文与训练开销

模型只在 16k NIAH 上训练，却被直接测试到 32k、64k 和 128k。结果如表 1。TTFT 加速从 $11.3\times$ 增至 $13.0\times$，显存缩减几乎恒定在 $3.42$–$3.43\times$。这组数字恰好印证方法结构：保留词元数不随上下文线性增长时，上层成本继续受控；但门层以下仍完整处理全部输入，因此峰值显存由未剪枝阶段主导，无法随长度继续扩大收益。

表 1（论文 Table 3）：Qwen3-1.7B 在 NIAH 不同测试长度上的外推结果

测试长度	TTFT 加速	峰值预填充显存缩减
16k	$11.3\times$	$3.42\times$
32k	$12.8\times$	$3.43\times$
64k	$13.0\times$	$3.43\times$
128k	$13.0\times$	$3.43\times$

训练成本没有因为双预算前向而整体翻倍。论文在单张 A100 上统计的均值是：普通 SFT 总训练时间 $2.2$ 小时后，TokenGate 为 $1.8$ 小时，比例 $0.8\times$；但单步时间从 $0.45$ 秒增至 $0.58$ 秒，比例 $1.3\times$。原因是 TokenGate 训练虽然每步更重，却在五个任务中的四个更早收敛。这张表同时说明“训练便宜”和“每步更慢”并不矛盾，比较口径必须分成总时长与单步开销。

表 2（论文 Table 2）：Qwen3-1.7B 在单张 A100 上的 SFT 与 TokenGate 训练时间

数据集	SFT 总时长	TokenGate 总时长	总时长比	SFT 单步	TokenGate 单步	单步比
RepoBench	1.7h	1.1h	$0.6\times$	0.50s	0.64s	$1.3\times$
2Wiki	3.8h	4.0h	$1.1\times$	0.34s	0.43s	$1.3\times$
MultiNews	1.3h	1.8h	$1.4\times$	0.21s	0.27s	$1.3\times$
BioDEX	1.2h	1.0h	$0.8\times$	0.23s	0.30s	$1.3\times$
NIAH	2.8h	1.0h	$0.4\times$	0.99s	1.24s	$1.3\times$
平均	2.2h	1.8h	$0.8\times$	0.45s	0.58s	$1.3\times$

4.3 消融：哪些设计不能省

图 7 左侧逐一移除全上下文损失、温度缩放和头级缩放。最明显的退化来自去掉全上下文项：模型在激进剪枝时看似仍能跑，但全够预算的精度本身下降，失去“训练后至少不伤原模型”的锚点。去掉温度或头级增益后，准确率—时延前沿也会内缩。这说明式 (6)、(7) 不是装饰性的超参数，而是让统一 checkpoint 跨预算工作所需的数值机制。

图 7 右侧比较门层 0、1、2。第 0 层明显 Pareto 次优，门层 1 与 2 则接近前沿。作者的解释是，第 0 层注意力后残差还没有足够语义信息，评分器难以判断与问题相关的实体；到第 1 层，表示已获得更丰富的上下文信号。这个结论也限定了方法的适用边界：门层位置必须调，不能把它当成任意插入都等价的正则项。

4.4 定性结果：保留什么，以及何时答错

论文对每个数据集检查真实样本，并在十等分位置上观察保留词元的分布。选择是全局的，没有退化为只看开头或只看结尾；但在 NIAH 中，$15{,}843$ 个词元的草堆里，位于约 $78\%$ 深度的 24 词元 needle 句有 21 个词元进入 $K=64$ 集合，模型直接答对。这是一个强证据：门控不是简单关键词匹配，它能把远距离且稀有的连续片段整体抬到高预算位置。

失败案例同样清楚。MultiNews 容易保留功能词而漏掉人名和地点，于是摘要事件方向正确但专有实体错误；BioDEX 在样例中只恢复两种药物反应中的一种。RepoBench-P 更反直觉，保留预算大量落在括号、逗号和代码结构符号上，作者推测上层把这些词元当作检索文件结构的锚点，而不是复制内容。表 3 汇总了这些模式，说明同一评分机制在不同任务上学出了不同策略。

表 3（论文 Table 4）：真实样例中的保留内容与失败方式

数据集	主要保留内容	观察到的失败
RepoBench-P	代码语法而非标识符，逗号与空参数列表突出，保留点较分散	小预算下生成错误代码行
2WikiMultihopQA	实体和日期，功能词较少	小预算混淆奖项；多跳答案需要较大预算
MultiNews	功能词与事件词，但人名、地点保留不足	事件大致正确，姓名、职业或地点错误
BioDEX	诊断与结果术语，以及列表连接词；数字保留很少	可能漏掉部分药物反应
NIAH-16k	needle 句作为连续片段保留，另有少量稀有子词	展示样例在全部预算下均精确

五、局限与适用边界

第一项结构性限制是显存收益在门层之上发生。门层以下仍要完整处理所有词元，峰值显存因此由完整前向主导；这也是表 1 中显存缩减稳定在约 $3.43\times$、没有随长度继续增加的原因。若未来把门放得更早，可以减少这段完整前向，但图 7 表明浅层评分信号不足，速度快与选择准之间存在真实张力。

第二项来自作者报告的语义遗漏。MultiNews 的摘要在低预算下会弄错名称或地点，BioDEX 会漏掉一项反应，2WikiMultihopQA 在多跳约束下需要更大预算才能与完整上下文一致。对高风险机器人场景，这种错误可能直接改变动作目标或安全属性，不能只按平均 F1 决定预算；应把关键实体、单位、限位和故障码加入不可剪枝集合，并在任务级最坏情况上验收。

第三项是实验覆盖。所有时间都是单张 A100、批大小为 1，主要报告预填充 TTFT 与峰值显存；论文没有给出服务端多租户吞吐、端到端生成时延、能耗或真实机器人闭环测试。现实部署还要考虑批处理利用率、CUDA Graph 形状、动态 $K$ 调度和多个请求共享设备的情况。作者在正文中特意说明未启用下层 KV 缓存剪枝，以隔离预填充收益；这既让比较干净，也意味着系统总收益仍取决于后续解码优化。

此外，仓库与论文的可复现性需要读者留意。论文正文印出的 GitHub 组织名在实际公开版本中并不是最终仓库地址，本轮核对的有效代码位于 token-gate-llms/TokenGate。公开代码包含训练、基准和 kernel 相关路径，但论文点到的大量数据准备、完整模型 checkpoint 与硬件条件仍决定复现成本。把“代码可见”当作“所有数字已独立复现”是不严谨的。

六、结论与对机器人系统设计的启示

TokenGate 的有效点不在于发明了更复杂的剪枝评分器，而在于把选择器、注意力和主干适配放进同一条梯度路径。这条路径绕开 top-$K$ 的离散边界：选择器决定哪些词元继续参与计算，门值则在被选中的词元之间继续改变注意力质量，因此任务损失能反馈到“该保留谁”。用一个 checkpoint 支持多个 $K$，再加上比 SFT 略短的总训练时间，使它具备成为现有 LLM 微调后处理组件的现实性。

对机器人基础模型，最值得验证的不是直接把语言基准数字搬进控制回路，而是三个受控问题：视觉词元的多尺度冗余是否比文本更明显；将安全关键区域、机器人状态和动作头词元列为 $\mathcal N$ 后，最坏情况性能是否可控；当相机帧率导致请求连续到达时，动态预算能否同时满足控制周期的硬截止时间。这些实验若成立，学习式预填充剪枝可能成为长观察历史与轻量动作模型之间的一层调度机制；若不成立，至少这篇论文提供了明确的测量范式与失败案例。

金句

TokenGate 没有让模型少看世界，而是让它先决定哪些输入值得进入更深的思考，再用可训练的门控把这个决定变成可优化的计算路径。