冻结多token预测加速Pixel上的Gemini Nano模型
阅读原文· research.google谷歌这篇技术博客值得端侧开发者细读,他们把多令牌预测硬是装进了已部署的 Nano 模型,Pixel 上生成加速五成,还省了 130MB 内存,零拷贝架构的想法挺巧,但没法直接复现,主要是开脑洞用的。
Google Research提出一种新架构,在已冻结的Gemini Nano v3模型上改造Multi-Token Prediction(MTP),以加速Pixel 9和10系列上的设备端推理。该方法基于EAGLE框架和CALM,无需单独训练占用内存的草稿模型,通过“晚期退出”策略实现加速。AI通知摘要和校对功能因此生成文本速度显著提升、能耗降低,开发者无需为每个新任务微调独立模型。
探索我们的众多重点领域
查看所有研究领域
应用 AI 与科学
地球 AI 健康 AI 科学 AI 可持续发展与危机韧性
基础 ML 与算法
算法与理论 信息检索 机器智能 机器感知 自然语言处理
人员、系统与量子 AI
人机交互与可视化 网络 量子 AI 负责任 AI 反滥用 软件工程 软件系统
了解更多
论文 项目
构建协作生态
数据集 获取高质量数据集以加速你的研究。工具与服务 探索我们最新的 AI 模型和产品。
开源 发现开源代码并与社区协作。
共同塑造未来
查看所有项目
教师项目 通过与大学教师的有意义互动,参与学术研究社区。学生项目 通过广泛的编程支持支持下一代研究者。
工作地点 在我们的全球办事处和研究实验室中找到你的位置。
将发现转化为现实世界的影响
人员 我们的研究者通过基础研究和应用研究推动计算机科学的进步。团队 协作小组应对世界上最具挑战性的 AI 问题。
研究
探索我们的众多重点领域
查看所有研究领域
应用 AI 与科学
基础 ML 与算法
人员、系统与量子 AI
了解更多
论文 项目
资源
构建协作生态
数据集 获取高质量数据集以加速你的研究。工具与服务 探索我们最新的 AI 模型和产品。开源 发现开源代码并与社区协作。
会议与活动
职业发展
共同塑造未来
查看所有项目
教师项目 通过与大学教师的有意义互动,参与学术研究社区。学生项目 通过广泛的编程支持支持下一代研究者。工作地点 在我们的全球办事处和研究实验室中找到你的位置。
博客
关于我们
将发现转化为现实世界的影响
人员 我们的研究者通过基础研究和应用研究推动计算机科学的进步。团队 协作小组应对世界上最具挑战性的 AI 问题。
Google Research
Google AI 了解我们所有的 AIGoogle DeepMind 探索 AI 的前沿Google Labs 尝试我们的 AI 实验
研究
资源
利用冻结的多Token预测加速Pixel上的Gemini Nano模型
2026年6月26日
Eden Cohen(研究产品经理)和 Michelle Ramanovich(研究经理),Google Platforms and Devices
我们介绍了一种方法,将多Token预测改造到冻结的生产模型上,从而加速设备端推理,无需单独的草稿模型带来的低效。
将强大的大语言模型(LLM)直接装进口袋,如今已通过Gemini Nano和Gemma等设备端模型成为现实。这项技术让手机上的日常功能成为可能——比如即时总结一大串通知,或校对重要短信——而且全程无需将你的私人数据发送到设备之外。但要让这些功能对日常用户真正有用,它们必须高效运行。
在移动设备上实现这样的速度是一项重大挑战。与庞大的服务器环境不同,手机在严格的能耗预算和硬性的内存(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)。 2. 验证:一个大模型(验证器)并行处理这些候选 token。如果候选 token 与大模型本应预测的结果一致,则接受它们;如果不一致,系统会回退到第一个分歧点。
然而,这会导致一些低效问题。运行一个独立的“草稿”模型(例如 128M 参数)会争抢有限的 RAM。此外,独立草稿模型“无视”主模型丰富的内部状态,仅基于文本历史来预测后续 token,而缺少主模型已经计算出的语义上下文。MTP 通过从独立架构转向集成架构来解决这些低效问题:我们不再训练一个单独的小语言模型来草拟 token,而是在主模型的最后几层上附加一个轻量化的 Transformer head,即 MTP head。
这种架构使用一个深层退出层进行草拟,从而利用了主模型主干网络已经完成的计算工作。MTP 头接收主模型最终的高维激活值(隐状态),并用它们来自回归地预测一系列未来的模型 token。
冻结主干网络的优势
虽然 MTP 头通常与主干网络一起预训练——例如在我们近期发布的 Gemma 4 模型中就是这样——但在利用已经部署在设备端的基础模型时,这种做法成本过高。相反,我们的工作重点是将草拟头改造为独立于预训练流水线运行的方式。
我们取一个完全训练好的 Gemini Nano v3 模型,冻结其权重,并在其最终层上附加一个密集的 Transformer 架构堆栈——即 MTP 头。我们仅训练这些参数,以最小化对未来模型 token 的预测误差。通过冻结主干网络,MTP 严格变成了一种效率优化,从而确保基础模型的能力或安全对齐不会退化。
由于错误的草拟会在验证过程中被丢弃,最终的输出与主模型保持逐比特一致,这使得我们能够在完全向后兼容的前提下推出效率更新。
零拷贝架构
虽然标准的 MTP 实现通过在主模型与草拟器之间共享静态参数(如嵌入向量权重)来优化训练效率,但设备端推理面临一个更严格的瓶颈:动态内存。即使共享了权重,如果草拟器独立处理上下文,它会产生并维护自己的键值缓存(KV cache),从而造成内存上的“双重负担”。考虑到移动设备内存有限,避免这种冗余至关重要。
为解决这一问题,我们设计了一种零拷贝架构,使得 MTP 头能够有效利用主模型的状态。MTP 头不是维护自己的历史记录,而是设计为直接交叉注意力到主模型的冻结键值缓存。这样一来,草拟器就能查询主干网络已经计算好的“记忆”和上下文,而无需重复计算。
这一设计带来了两方面的效率提升。首先,它消除了草拟器的预填充延迟:通过利用现有的缓存,MTP 头无需额外时间来处理提示词。其次,它减少了运行时的内存占用。与独立草拟器相比,我们观察到每个实例节省了 130MB 内存,这得益于省去了草拟器的嵌入向量查找表、预填充点积注意力变体以及特定应用调优参数。
通过利用主模型的隐藏状态和 KV 缓存,MTP 头部生成候选 token,由主干网络并行验证,从而消除冗余的预填充延迟,并将内存使用量降低最多 130MB。
解锁更丰富的表示
在我们的实验中,我们发现 MTP 起草器始终能产生更准确的 token 预测,与参数数量相当且“独立运行的起草器”相比,在 Pixel 9 设备上实现了 50% 或更高的速度提升[aef552],具体取决于任务。
这种性能差距源于 MTP 能够访问更丰富的表示。与将主模型视为黑盒的独立起草器不同,MTP 头部直接利用已由更大主干网络处理过的最终激活值:
_指令遵循:_ 在摘要或带复杂约束的重写等任务中,MTP 显著优于独立微调的起草器。
_可预测的文本结构:_ 对于结构可预测性高的任务(例如智能回复),MTP 头部有效学习了主模型的句法模式,实现 token 接受率最高提升 55%。
实际影响
为了在 Pixel 9 和 Pixel 10 设备上部署 MTP,我们重新设计了设备端推理栈,以处理验证阶段与起草阶段之间的复杂依赖关系。
结果验证了架构选择的有效性。在生产工作负载中,例如 AI 通知摘要和校对,MTP 每次推理平均能正确预测接近两个额外的 token。此外,更少的验证步骤意味着唤醒重型处理器的时间更少,从而降低能耗并延长电池续航。
_MTP 与各应用专用独立调优起草器在多种 Pixel 9 应用中对 Gemini Nano token 生成的影响对比。_
未来方向
我们期待将 MTP 集成到未来的 Pixel 设备上,同时探索替代架构——包括并行解码和无辅助头部的范式——以进一步降低草案生成延迟,并在严格的移动端约束下提高同时 token 验证的数量。
我们还在研究更高效地处理语言生成内在模糊性的方法。虽然标准推测解码假设存在单一最佳未来…