在Pixel上加速Gemini Nano模型的冻结多Token预测技术

将强大的大语言模型装入口袋，对于Gemini Nano和Gemma等设备端模型而言已成现实。这项技术让手机上的日常功能得以实现——例如即时汇总大量通知或对重要短信进行语法校对——且无需将私人数据上传至云端。然而，要让这些功能真正服务于日常用户，就必须保证极高的运行效率。

在移动设备上实现这种速度是一项重大挑战。与庞大的服务器环境不同，手机在能耗预算和内存（RAM）方面都受到严格限制。此外，标准语言模型以"自回归"方式生成文本——即每次仅处理并输出一个词（或Token）。这种逐步生成的方式造成了瓶颈，在占用内存带宽的同时未能充分利用手机的处理能力，最终可能拖慢用户体验并加速电池耗尽。

为突破这一瓶颈，谷歌宣布推出一种新架构，将多Token预测（MTP）技术移植到现有的"冻结"Gemini Nano v3模型上。在EAGLE框架和置信自适应语言建模（CALM）等前期方案的基础上，团队专门针对移动端环境设计了新的架构组件，以最大化效率收益。此前的公告中已介绍了通过MTP加速Gemma 4并向开发者开放的相关进展。

本文聚焦于边缘计算在极端约束条件下的独特挑战。该方案已在Pixel 9和10系列上线，作为开箱即用的加速方案存在。对用户而言，AI通知摘要和语法校对等功能将以更快的速度、更低的能耗生成文本；对开发者而言，则省去了一大痛点：无需为每项新任务单独微调占用大量内存的草稿模型，即可实现高速的设备端AI推理。

MTP的技术演进

MTP是在推测解码技术的基础上发展而来的。在传统方案中，生成N个Token需要对大模型进行N次前向推理。推测解码将这一过程解耦为两个阶段：由小型草稿模型快速生成候选Token序列，再由主模型并行验证。

然而，这种方式存在一定的低效之处。独立运行的"单机版"草稿模型（例如1.28亿参数）会与有限的RAM资源形成竞争。此外，独立草稿模型对主模型的内部状态"一无所知"，仅依据文本历史预测下一个Token，无法利用主模型已计算的语义上下文。MTP通过将架构从独立式转变为集成式来解决上述问题。它不再单独训练一个小型语言模型来起草Token，而是在主模型的最终层后附加一个轻量级Transformer头，即MTP头。

这一架构采用深层退出层进行草稿生成，充分利用主模型主干已完成的计算工作。MTP头获取主模型最终的高维激活值（隐藏状态），并以自回归方式预测一系列未来Token。

冻结主干的设计意义

尽管MTP头通常与主干网络同步预训练——如近期发布的Gemma 4模型——但对于已在设备端部署的基础模型，这种做法并不可行。因此，该研究工作的重点在于将草稿头以独立于预训练流程的方式进行移植。

具体做法是：取一个完整训练好的Gemini Nano v3模型，冻结其权重，然后在最终层后附加一个密集Transformer堆栈，即MTP头，并仅训练这部分参数，使其在预测未来Token时的误差最小化。在主干冻结的前提下，MTP严格作为效率优化手段存在，确保基础模型的能力和安全对齐不会因此退化。

由于验证阶段会丢弃错误的草稿，最终输出与主模型保持逐位一致，从而实现效率更新的完全向后兼容。

零拷贝架构：解决内存瓶颈

标准MTP实现通常通过共享静态参数（如嵌入权重）来提升训练效率，而设备端推理面临的更严峻瓶颈是动态内存。即便权重共享，若草稿模型独立处理上下文，也会因生成并维护自身的键值（KV）缓存而造成"双重内存负担"。鉴于移动端内存极为有限，避免这种冗余至关重要。

为此，团队设计了一种零拷贝架构，使MTP头能有效复用主模型的状态。MTP头不再维护自身的历史记录，而是直接交叉注意主模型冻结的KV缓存，从而在不产生数据复制的情况下查询主干已计算的"记忆"与上下文。

这一设计带来两项效率收益：其一，消除草稿模型的预填充延迟——通过复用已有缓存，MTP头无需额外时间处理提示词；其二，降低运行时内存占用。与独立草稿模型相比，每个实例节省了约130MB内存，节省来源包括草稿嵌入查找表、预填充点注意力变体以及应用专属调优参数。

通过复用主模型的隐藏状态和KV缓存，MTP头生成的候选Token由主干并行验证，在消除冗余预填充延迟的同时，内存用量最高可减少130MB。

实验结果

实验表明，与参数量相当的"独立草稿模型"相比，MTP草稿模型在Pixel 9设备上的Token预测准确率持续更高，在不同任务中实现了50%甚至更高的速度提升。

这一性能差距源于MTP能够获取更丰富的表征。与将主模型视为黑盒的独立草稿模型不同，MTP头直接利用由更大主干处理后的最终激活值，从而做出更准确的预测。

在实际生产负载中，如AI通知摘要和语法校对，MTP平均每次推理可正确预测近两个额外Token。更少的验证步骤意味着重型处理器的唤醒时间更短，从而降低能耗、延长电池续航。

未来展望

团队期待在未来的Pixel设备上持续集成MTP技术，并探索替代架构——包括并行解码以及无辅助头的新范式——以在严格的移动端约束下进一步降低草稿延迟并提升同步Token验证数量。

团队还在研究如何更高效地处理语言生成中固有的模糊性问题。标准推测解码假设存在唯一最优的未来路径，而研究人员正在开发能够并行探索多个分支可能性的技术，旨在即使在不确定的上下文中也能最大化接受长序列的概率。此外，团队还在研究验证宽松化：对于特定应用场景，放宽草稿与验证之间严格的Token精确匹配要求，以进一步提升边缘端效率。

致谢

本项工作是设备端大语言模型效率优化系列研究的组成部分，参与成员包括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团队在审阅、反馈与支持方面的宝贵贡献。

Q&A

Q1：多Token预测（MTP）是什么技术？它如何提升手机上的AI速度？

A：多Token预测（MTP）是一种在主模型最终层后附加轻量级Transformer头的技术。它利用主模型已计算的隐藏状态和KV缓存，一次性预测多个未来Token，再由主模型并行验证。这样既避免了逐个Token生成的瓶颈，又无需额外训练独立草稿模型，在Pixel 9设备上可实现50%以上的推理速度提升。

Q2：MTP技术为什么要"冻结"Gemini Nano v3模型的权重？

A：冻结主模型权重是为了在不影响已部署模型性能和安全对齐的前提下，将MTP作为纯粹的效率优化手段叠加进去。只训练新附加的MTP头参数，确保最终输出与原始模型逐位一致，从而实现完全向后兼容，可以安全地推送效率更新而不破坏已有功能。

Q3：MTP技术是如何解决手机内存不足问题的？

A：团队设计了零拷贝架构，让MTP头直接交叉注意主模型已有的KV缓存，而非自行维护一套缓存。这消除了草稿模型的预填充延迟，并使每个实例的运行时内存占用减少约130MB，有效解决了移动端内存有限的瓶颈问题。