斯坦福和UC伯克利研究员发布 TokenGate，一个用于长文本预填的模块。 TokenGate 学习到的令牌浏览让现有的 LLM 省略其输入。 只有那些重要的东西才能完全处理研究组报告1.4x至11.3x较快的时速为第一键(TTFT)，没有失精。 峰值预填内存下降1.5x至3.4x. 密码和检查站位于Apache-2.0下的GitHub。
TL；数据
- MLP门：一个小型MLP门和LoRA(r=16)，加入Qwen3-1.7B或Qwen3.5-2B. 对上下文进行了最多128k令牌的测试。
- 运行在：纸质计时使用1A100，批量尺寸为1. 培训使用了4个A100。
- 业绩：对照6个基线，对5项任务进行优化。
- 最佳：NIAH 16k在11.3x更快的TTFT和3.42x更低的内存。 TTFT以128k获得达到 13.0x.
- 最差：一些任务速度降至1.4x。
托肯盖特是什么？
TokenGate是一个滴入组件，新增可列车代币普鲁因预训LLMs. 变形器将每个符号推穿每个层。 注意力成本随长度增长四倍。 免训练的修炼成本低廉，但以高修炼率降低精度。 建筑改造往往需要从零开始的预训。
TokenGate是怎么学到代币的？
前传分为三个阶段：
- 标准前方：门下L层下层照常处理所有N符号。
- 门模块：一个MLP从L层的留守后每个符号中得分。
- 向前冲：L 上方的层只处理这些 K 令牌，绝不为已冲动的 KV 条目创建。
只有提示的文档部分是可读的。 问题从未解决。
什么是象征性的关注？
Top-K选择不相异，这阻断了端到端的训练。 TokenGate借用了Mixture-of-Experts路由器的技巧。 每个记号的得分都成了大门上方的注意力记录的偏差。 这种偏差让梯度返回计分员的路径。
可学习的人均收益和可学习的温度可以完善这一点。 显示键数较少的每个查询的温度校正。 训练用的是直接的下手每一步骤平均损失为全文损失和随机预算损失。 因此，一个训练有素的大门是跨预算的。
TokenGate在基准测试上表现如何？
研究团队在Repo Bench-P,2Wiki MultihopQA,MultiNews,BioDEX和针入海斯塔克(NIAH)上测试了Qwen3-1.7B. 基线为FastKV、SlimInfer、DASH、GemFilter、FlexPrefilling和MInform，位于同样的LoRA主干线上。
- Versus完全预装：1.4x至11.3x更快的TTFT和1.5x至3.4x更低的内存。
- Versus，次佳基线：1.2x至6.4x更快的TTFT，最多2.1x更低的内存。
- 2Wiki MultihopQA：以全文本精度的95%，预填运行速度为3.0x，内存减少2.7x.
- 较长的背景：经过16k训练的大门在32k达到12.8x，在64k和128k达到13.0x。 记忆储蓄停留在3.43x附近。
训练有素的选手在同一个层次上战胜了FastKV的热情。 甚至随机K选手在TokenGate主干线上的分数更高。 研究团队对骨干协同改造给予肯定。
它是否在混合关注模式上起作用？
是的，Qwen3.5-2B混合了Gated DeltaNet和全注层。 在NIAH 32k上，TokenGate在第3层与全预装精度匹配。 它跑的快于满载预装3.8x，快于FastKV. 第1层或第2层的测量达到7.8x，精确度在6%以内。 FastKV只能使用prune关注层。
TokenGate vs FastKV, SlimInfer, GemFilter 和 DASH 游戏游戏
自报收益采用不同的设置，因此不具有直接可比性。 螺旋层来自TokenGate的附录C.2.
你怎么部署托肯盖特？
GitHub Repo需要Python 3.10+和CUDA PyTorch 2.9+. 可选的FlashAtthy-2密钥-比斯内核加速稠密的Qwen3. 它花费1.0x到1.1x 一个普通的呼叫。 没有它，推论就回到FlexAttenty.
{状态}占位符标记了可保存的文档。 生成需要最大 kept prefill tokens 预算。 检查站负责Qwen3-1.7B上的所有5项任务，外加Qwen3.5-2B NIAH门。 示例：liana313/tokengate-niah-qwen3-1.7b.
培训反映了TRL/PEFT的配方。 每一步比平原LoRA慢1.3x. 5项任务中的4项任务运行速度更快。 平均培训时间为1.8小时，而为2.2个A100小时。
关键外卖
- TokenGate 新增通过一个小门和LoRA向现有的LLMs学习了代币追逐。
- Prefill运行1.4x至11.3x运行速度更快，1.5x至3.4x的峰值存储器降低。
——托肯格特式关注使上K型选择式可列车，MOE-router型。
- 它比TTFT上最优秀的推力基线要高6.4倍
- 它工作在杂交的Qwen3.5-2B，其中FastKV不能prune DeltaNet层。
请检查access-date=中的日期值 (帮助) 页面存档备份，存于互联网档案馆页面存档备份，存于互联网档案馆 GitHub repo and Hugging Face纸页。 并加入我们的通讯。 等等你在电报里吗现在你也可以加入我们的电报。
需要和我们合作推广您的 GitHub Repo OR Hugging Face Page OR Product Release OR Webinar 等吗？联系我们
Asif Razzaq是Marktechpost AI Media Inc.的首席执行官。 作为一个有远见的企业家和工程师，阿西夫致力于利用人工智能的潜力促进社会福祉。 他最近的努力是人工智能媒体平台的发布，该平台以深入报道机器学习和深层学习新闻而著称，这些新闻在技术上既可靠又为广大观众所容易理解。 该平台每月有超过200万次观看，表明其受观众欢迎。
本文由 AI图鉴自动抓取并翻译，版权归原作者所有；如需查看原文请点下方链接。
Stanford and UC Berkeley researchers released TokenGate, a module for long-context prefill. TokenGate learned token pruning lets an existing LLM skim its input. Only the tokens that matter get full processing. The research team reports 1.4x to 11.3x faster time-to-first-token (TTFT) with no accuracy loss. Peak prefill memory drops 1.5x to 3.4x. Code and checkpoints are on GitHub under Apache-2.0.
TL;DR
- MLP gate: A small MLP gate and LoRA (r=16) added to Qwen3-1.7B or Qwen3.5-2B. Tested on contexts up to 128k tokens.
- Runs on: Paper timings use 1 A100 at batch size 1. Training used 4 A100s.
- Performance: Pareto-optimal on 5 tasks against 6 baselines.
- Best: NIAH 16k at 11.3x faster TTFT and 3.42x lower memory. TTFT gain hits 13.0x at 128k.
- Worst: Speedup drops to 1.4x on some tasks.
What is TokenGate?
TokenGate is a drop-in component that adds trainable token pruning to pretrained LLMs. Transformers push every token through every layer. Attention cost grows quadratically with length. Training-free pruning is cheap but loses accuracy at high pruning rates. Architectural changes often need pre-training from scratch.
How does TokenGate learned token pruning work?
The forward pass splits into 3 phases:
- Standard forward: layers below gate layer L process all N tokens as usual.
- Gate module: an MLP scores each token from its post-attention residual at layer L. A top-K selector keeps the highest-scoring tokens.
- Pruned forward: layers above L process only those K tokens and never create KV entries for pruned ones.
Only the document part of the prompt is prunable. The question is never pruned.
What is token-gated attention?
Top-K selection is not differentiable, which blocks end-to-end training. TokenGate borrows the Mixture-of-Experts router trick. Each kept token’s score becomes a bias on attention logits above the gate. That bias gives gradients a path back to the scorer.
A learnable per-head gain and a learnable temperature refine this. The temperature corrects for each query seeing fewer keys. Training uses plain next-token loss. Each step averages a full-context loss and a random-budget loss. So 1 trained gate works across budgets.
How does TokenGate perform on benchmarks?
The research team tested Qwen3-1.7B on RepoBench-P, 2WikiMultihopQA, MultiNews, BioDEX and needle-in-a-haystack (NIAH). Baselines were FastKV, SlimInfer, DASH, GemFilter, FlexPrefill and MInference, on the same LoRA backbone.
- Versus full prefill: 1.4x to 11.3x faster TTFT and 1.5x to 3.4x lower memory.
- Versus the next-best baseline: 1.2x to 6.4x faster TTFT and up to 2.1x lower memory.
- 2WikiMultihopQA: at 95% of full-context accuracy, prefill runs 3.0x faster with 2.7x less memory.
- Longer contexts: a gate trained at 16k reaches 12.8x at 32k and 13.0x at 64k and 128k. Memory savings stay near 3.43x.
The learned selector beats FastKV’s heuristic at the same layer. Even random-K selection scores higher on the TokenGate backbone. The research team credits backbone co-adaptation.
Does it work on hybrid-attention models?
Yes, Qwen3.5-2B mixes Gated DeltaNet and full-attention layers. On NIAH 32k, TokenGate at layer 3 matched full-prefill accuracy. It ran 3.8x faster than full prefill and 2.1x faster than FastKV. Gating at layers 1 or 2 reached 7.8x, within 6% accuracy. FastKV can only prune attention layers.
TokenGate vs FastKV, SlimInfer, GemFilter and DASH
Self-reported gains use different setups, so they are not directly comparable. Pruning layers come from TokenGate’s Appendix C.2.
How do you deploy TokenGate?
The GitHub Repo needs Python 3.10+ and CUDA PyTorch 2.9+. An optional FlashAttention-2 key-bias kernel speeds up dense Qwen3. It costs 1.0x to 1.1x a plain call. Without it, inference falls back to FlexAttention.
A {state} placeholder marks the prunable document. generate takes a max_kept_prefill_tokens budget. Checkpoints cover all 5 tasks on Qwen3-1.7B, plus a Qwen3.5-2B NIAH gate. Example: liana313/tokengate-niah-qwen3-1.7b.
Training mirrors a TRL/PEFT recipe. Each step is 1.3x slower than plain LoRA. Runs converged faster on 4 of 5 tasks. Mean training time was 1.8 versus 2.2 A100 hours.
Key Takeaways
- TokenGate adds learned token pruning to existing LLMs via a small gate and LoRA.
- Prefill runs 1.4x to 11.3x faster with 1.5x to 3.4x lower peak memory.
- Token-gated attention makes top-K selection trainable, MoE-router style.
- It beats the next-best pruning baseline by up to 6.4x on TTFT.
- It works on hybrid Qwen3.5-2B, where FastKV cannot prune DeltaNet layers.
Check out the Paper, GitHub repo and Hugging Face paper page. Also, feel free to follow us on Twitter and don’t forget to join our 150k+ML SubReddit and Subscribe to our Newsletter. Wait! are you on telegram? now you can join us on telegram as well.
Need to partner with us for promoting your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar etc.? Connect with us
Asif Razzaq is the CEO of Marktechpost AI Media Inc.. As a visionary entrepreneur and engineer, Asif is committed to harnessing the potential of Artificial Intelligence for social good. His most recent endeavor is the launch of an Artificial Intelligence Media Platform, Marktechpost, which stands out for its in-depth coverage of machine learning and deep learning news that is both technically sound and easily understandable by a wide audience. The platform boasts of over 2 million monthly views, illustrating its popularity among audiences.
Auto-fetched and translated by AI Atlas. All rights belong to the original author.