TokenGate：可学习 token 剪枝加速长上下文
AI 总结
有研究者在 r/LocalLLaMA 发布论文，介绍 TokenGate——一个小型即插即用组件，可扩展已训练好的模型，在长上下文输入上实现端到端可学习的 token 剪枝，从而加快预填充（prefill）并降低内存占用。为绕开 top-K 算子不可微的问题，作者提出 token-gated attention，借鉴 Mixture-of-Experts 思路，用 token 重要性分数重新加权注意力 logits。作者自报结果称，预填充阶段首 token 延迟（TTFT）加快 1.4–11.3 倍、峰值内存降低 1.5–3.4 倍且精度不降；相比此前最佳的 token 剪枝方法，预填充最高快 6.4 倍、峰值内存低 2 倍。作者称 TokenGate 可通过轻量适配原生扩展到 Qwen3 等标准 Transformer 以及 Qwen3.5 等混合注意力模型，代码已在 GitHub 公开。
深度见解
作者指出语言模型会让每个输入 token 经过每一层，长上下文预填充成本高昂，并将 TokenGate 定位为在该阶段剪掉不重要 token 的即插即用方案。
| Expert human readers process large inputs (e.g., code repositories and long documents) by skimming to closely read only what matters. On the other hand, language models process all input tokens through each model layer, resulting in significant processing costs for long-context inputs. We introduce TokenGate, a small, drop-in component that extends existing trained models to enable end-to-end learnable token pruning over long-context inputs for fast and memory-efficient prefill. To enable end-to-end learning despite the non-differentiable top-K operator for token pruning, we introduce token-gated attention, a mechanism inspired by Mixture-of-Experts that re-weights attention logits using token importance scores. Across diverse tasks, TokenGate achieves 1.4-11.3× faster time-to-first token (TTFT) and 1.5-3.4× lower peak memory at no accuracy degradation for prefill in transformers. Compared to the best-performing prior methods for token-pruning, TokenGate gives up to 6.4× faster prefill and 2× lower peak memory. TokenGate natively extends both standard transformers (e.g., Qwen3) and hybrid-attention models (e.g., Qwen3.5) with simple lightweight adaptation comparable to standard parameter-efficient fine-tuning. We make our code available at: this https URL. submitted by /u/pmttyji [link] [comments] |