TokenGate：学习式筛选，降低预填充开销

1. 让现有模型学习如何筛选词元

这种方法让模型选择要处理的部分，就像人浏览文档时会仔细阅读重要段落。现有模型通常在每一层计算所有词元。TokenGate 在指定层为每个文档词元评估重要性，只保留设定数量并送往更上层。未保留的词元不参与后续层计算，也不进入这些层生成的键值缓存，即存储并复用先前计算结果的内存。

现有词元筛选方法无需训练，而是依据注意力分数等规则筛选词元，这类分数表示词元之间的关联程度。论文指出，仅使用早期层分数的方法在大量削减词元时会损失准确率。TokenGate 通过在现有模型上继续训练，对模型和筛选器进行联合微调，学习应保留哪些词元。

2. 联合微调筛选器与模型

训练时，词元筛选分数也会影响上层的注意力计算，让模型学习哪些选择有助于给出正确答案。由于同时纳入完整上下文与缩减上下文的训练结果，训练一次的筛选器可用于多种词元预算。

仓库中的推理示例通过标记文档区间并指定保留词元数来运行。以下相关片段原样摘自 README 代码。

for k in (None, 2000, 500, 100):          # None = keep every state token
    out = model.generate([ids], tok, max_kept_prefill_tokens=k)
    print(k, f"{out.prefill_s * 1e3:5.0f} ms", out.answers[0])

3. 长上下文中延迟降幅扩大

论文使用 16k 长度的大海捞针任务训练 Qwen3-1.7B，并在更长上下文中评估。大海捞针任务要求找出埋在长文档中的特定事实并作答。

下表展示了与处理完整上下文的微调模型达到相同准确率时的结果，比较了处理输入到输出首个词元的时间（TTFT）以及峰值内存用量。

表：TokenGate 在不同上下文长度下的节省效果

条件：Qwen3-1.7B，bf16，以 16k 大海捞针任务训练，批大小 1，1 张 A100，相对于完整上下文

单位：倍

上下文越长，加速幅度越大，而内存节省幅度维持在 3.42~3.43 倍。论文分析认为，这是因为门控下方的层仍处理所有词元，影响了峰值内存。

4. 仓库测试与任务间差异

仓库公布的独立基准测试，即用于比较性能的标准测试，使用 1 张 H100、批大小 1，对训练时留出的测试数据进行测量。下表中的原始模型指训练前的模型，不剪枝则是针对任务微调门控与模型后保留所有词元的准确率基线。

K=2000 表示保留 2000 个文档词元。代码补全分数衡量与正确代码的相似程度，药物反应抽取分数衡量抽取的完整性和准确性。

表：仓库公开测试结果

条件：基于 Qwen3-1.7B，1 张 H100，批大小 1，与训练数据分离的测试数据，K=2000

单位：TTFT 为相对于训练前原始模型的倍数

代码补全的分数低于不剪枝结果，但药物反应抽取略高。仓库的另一个大海捞针演示中，在约 32k 的文档中仅保留 100 个词元，五个案例仍全部找到答案，TTFT 加速为 11.18 倍。这是在五份文档、针的位置为 11~100%、使用 1 张 GPU 条件下的结果。

5. 速度、内存收益与代价

论文报告，在主要实验中，与完整上下文微调模型达到相同准确率时，预填充加速了 1.4~11.3 倍，峰值内存节省了 1.5~3.4 倍。这是在 Qwen3-1.7B、bf16、1 张 A100、批大小 1 条件下评估多个任务的结果。论文第 5.1 节指出，相比表现次优的筛选方法，加速最高为 6.4 倍，内存节省最高为 2.1 倍。

节省推理成本的同时，也需要付出训练筛选器和准备运行环境的成本。论文附录中使用 Qwen3-1.7B、1 张 A100 的实验显示，单个训练步骤比标准微调慢 1.3 倍，但五个任务的平均总训练时间为标准方式的 0.8 倍。

为实现默认配置下的快速运行，建议构建 CUDA 内核；推理也可以不使用内核，通过替代路径运行。训练时若使用替代路径，则需显式选择。

准确率损失也会随任务和词元预算而出现，论文中的摘要案例正确概括了事件，却弄错了姓名和地点。

TokenGate 为减少长上下文预填充开销提供了实用选择，但收益取决于训练数据以及针对任务的词元筛选。