冻结多token预测加速Pixel上的Gemini Nano模型
Accelerating Gemini Nano models on Pixel with frozen Multi-Token Prediction
Google Research提出一种新架构，在已冻结的Gemini Nano v3模型上改造Multi-Token Prediction（MTP），以加速Pixel 9和10系列上的设备端推理。该方法基于EAGLE框架和CALM，无需单独训练占用内存的草稿模型，通过“晚期退出”策略实现加速。AI通知摘要和校对功能因此生成文本速度显著提升、能耗降低，开发者无需为每个新任务微调独立模型。
谷歌这篇技术博客值得端侧开发者细读，他们把多令牌预测硬是装进了已部署的 Nano 模型，Pixel 上生成加速五成，还省了 130MB 内存，零拷贝架构的想法挺巧，但没法直接复现，主要是开脑洞用的。
我们提出了一种方法，可将多 token 预测改造到已冻结的生产模型上，从而加速端侧推理，同时避免独立草稿模型带来的低效。
有了 Gemini Nano 和 Gemma 这类端侧模型，把强大的大语言模型（LLM）装进口袋如今已成为现实。这项技术让手机上的日常功能成为可能——例如即时汇总一连串通知，或校对一条重要的短信——而且全程无需将你的隐私数据发送到设备之外。但要让这些功能对普通用户真正有用，它们必须运行得非常高效。
在移动设备上实现这种速度是一项重大挑战。与庞大的服务器环境不同，手机在严格的能耗预算和硬性内存（RAM）限制下运行。此外，标准语言模型以“自回归”方式生成文本——也就是说，它们一次只处理和输出一个词（或 token）。这种逐步处理的方式形成了瓶颈，既未能充分利用手机的处理能力，又给其内存带宽带来压力，最终可能拖慢用户体验并消耗电池电量。
为克服这一瓶颈，我们宣布了一种新架构，可将多 token 预测（MTP）改造到现有的、“冻结”的 Gemini Nano v3 模型上。在 EAGLE 框架和 Confident Adaptive Language Modeling（CALM）等先前方法的基础上，我们设计了新的架构组件，专门针对移动环境最大化这些效率收益。我们近期的公告重点介绍了用 MTP 加速 Gemma 4，并将其开放给开发者使用。
今天的文章探讨了边缘计算所面临的独特而极端的约束。这一方法近期已在 Pixel 9 和 10 系列上推出，可作为一种开箱即用的加速手段。对用户而言，这意味着 AI 通知摘要和校对等功能生成文本的速度显著更快，能耗也更低。对开发者而言，它消除了一大痛点：无需为每一项新任务微调单独的、占用大量内存的草稿模型，即可交付高速的端侧 AI。
一种“延迟退出”策略
MTP 建立在投机解码的演进之上。在传统设置中，生成 N 个 token 需要对大模型进行 N 次前向传播。投机解码将这一过程解耦为两部分：
- 草稿：一个更小、更快的近似模型（即“草稿模型”）生成一小段候选 token 序列（例如 3 个 token）。
- 验证：一个大模型（即“验证器”）并行处理这些候选 token。如果候选 token 与大模型原本会预测的结果一致，则被接受。如果不一致，系统则回滚到第一个出现分歧的位置。
然而，这会带来一些效率上的问题。运行一个单独的“独立”草稿模型（例如 128M 参数）会争抢有限的内存。此外，独立草稿模型对主模型丰富的内部状态是“盲”的，它仅根据文本历史预测后续 token，而无法利用主模型已经计算出的语义上下文。MTP 通过从独立架构转向集成架构，解决了这些效率问题。我们不再训练一个单独的小型语言模型来起草 token，而是在主模型的最后几层附加一个轻量级的 Transformer 头，即 MTP 头。
这种使用深层出口层进行起草的架构，利用了主模型主干已经完成的工作。MTP 头接收主模型最终的高维激活值（隐藏状态），并利用它们自回归地预测未来 token 序列。
冻结主干带来的优势
虽然 MTP 头通常与主干一起预训练——例如我们最近发布的 Gemma 4 模型——但在利用已部署的端侧基础模型时，这种做法是行不通的。相反，我们的工作重点是对草稿头进行改造，使其独立于预训练流程运行。
我们取一个完全训练好的 Gemini Nano v3 模型，冻结其权重，并在最后几层附加一个密集的 Transformer 堆栈——即 MTP 头。我们只训练这些参数，以最小化对未来 token 的预测误差。在主干冻结的情况下，MTP 严格来说成为一种效率优化，确保基础模型的能力或安全对齐不会退化。
由于错误的草稿在验证阶段会被丢弃，最终输出与主模型保持逐比特完全一致，这使我们能够在完全向后兼容的前提下推出效率更新。
零拷贝架构
虽然标准 MTP 实现通过在主模型与草稿模型之间共享静态参数（如嵌入向量权重）来优化训练效率，但端侧推理面临更严苛的瓶颈：动态内存。即便权重共享，如果草稿模型独立处理上下文，它仍会因生成并维护自己的键值（KV）缓存而对内存造成"双重负担"。考虑到移动设备内存有限，避免这种冗余至关重要。
为解决这一问题，我们设计了一种零拷贝架构，让 MTP 头能够有效利用主模型的状态。MTP 头不再维护自己的历史，而是被设计为直接交叉注意力到主模型冻结的 KV 缓存。这使得草稿模型能够查询主干网络已计算好的"记忆"和上下文，而无需重复计算。
这一设计带来两项效率提升。首先，它消除了草稿模型的预填充延迟：通过利用现有缓存，该头无需额外时间来处理提示词。其次，它降低了运行时内存占用。与独立草稿模型相比，我们观察到每个实例节省了 130MB，这得益于省去了草稿模型的嵌入向量查找表、预填充点积注意力变体以及应用特定的调优参数。
通过利用主模型的隐藏状态和 KV cache，MTP head 生成候选 token，由主干网络并行验证，从而消除了冗余的 prefill 延迟，并将内存占用最多减少 130MB。
释放更丰富的表征
在我们的实验中，我们发现 MTP drafter 始终能产生更准确的 token 预测，从而在 Pixel 9 设备上实现 50% 或更高的加速[aef552]，具体取决于任务，与参数量相当的“独立 drafter”相比。
这一性能差距源于 MTP 能够访问更丰富的表征。与将主模型视为黑盒的独立 drafter 不同，MTP head 直接利用已经由更大主干网络处理过的最终激活值：
- 指令遵循：在带有复杂约束的摘要或改写等任务中，MTP 显著优于独立微调的 drafter。
- 可预测的文本结构：对于结构可预测性高的任务（例如智能回复），MTP head 有效学习了主模型的句法模式，token 接受率提升最高达 55%。
实际影响
为了在 Pixel 9 和 10 设备上部署 MTP，我们重新设计了端侧推理栈，以处理验证阶段与起草阶段之间的复杂依赖关系。
结果验证了这些架构选择。在 AI Notification Summaries 和 Proofread 等生产工作负载中，MTP 每次推理过程平均能正确预测近两个额外的 token。此外，更少的验证步骤意味着唤醒重型处理器的次数更少，从而降低能耗并延长电池续航。
在多种 Pixel 9 应用中，MTP 与针对特定应用单独调优的 drafter 相比对 Gemini Nano token 生成的影响。
未来方向
我们期待在未来的 Pixel 设备上集成 MTP，并探索替代架构——包括并行解码和无辅助头的范式——以进一步降低 draft 延迟，并在严格的移动端约束下提高同时验证的 token 数量。
我们还在研究如何更高效地处理语言生成固有的歧义性。标准的投机解码假设存在单一最优未来路径，而我们正在开发允许模型并行探索分支可能性的技术。这旨在即使在不确定的上下文中也能最大化接受长序列的概率。此外，我们正在研究验证宽松度：针对特定用例放宽 draft 与验证之间严格的精确 token 匹配要求，从而为边缘端带来进一步的效率提升。
致谢
这项工作是我们优化端侧 LLM 效率努力的一部分，参与人员包括 Filippo Galgani、Omri Homburger、Pooja Consul、Matthew Markwell 和 Vivek Kumar。部分内容基于 Google DeepMind 中 Gemini 团队的开发成果：Tal Schuster、Ziwei ji、Ivan Korotkov 和 Ganesh Jawahar。我们还要衷心感谢 Nadav Bar、Utku Evci、Nir Shabat、Joe Zou，以及 Google Research、Google Deepmind 和 Platforms & Devices 各团队提供的评审、宝贵反馈与支持。
- 与 MTP 更新前的 Pixel 9 和 Pixel 10 手机相比。
来源：Google Research：Blog（网页） · research.google