利用冻结的多令牌预测在Pixel上加速Gemini Nano模型（10分钟阅读）

# 在 Pixel 上使用冻结的多 Token 预测加速 Gemini Nano 模型 来源：https://research.google/blog/accelerating-gemini-nano-models-on-pixel-with-frozen-multi-token-prediction/ 如今，像Gemini Nano (https://developer.android.com/ai/gemini-nano)和Gemma (https://deepmind.google/models/gemma/)这样的设备端模型，让强大的大型语言模型 \(LLMs\) 直接放在口袋里成为现实。这项技术为手机带来了日常功能——例如即时总结一堆通知或校对重要短信——所有这些都无需将你的私人数据发送到设备之外。但要让这些功能对日常用户真正有用，它们必须以极高的效率运行。 在移动设备上实现这种速度是一项重大挑战。与庞大的服务器环境不同，手机在严格的能耗预算和有限的内存 \(RAM\) 限制下运行。此外，标准语言模型是“自回归式”生成文本的——意味着它们一次只处理并输出一个单词（或 token）。这种逐步处理的过程造成了瓶颈，使得手机的处理能力未被充分利用，同时加剧了内存带宽的压力，最终可能拖慢用户体验并耗尽电池。 为了克服这个瓶颈，我们宣布了一种新架构，它将多 Token 预测 \(MTP\) 改造到现有的、“冻结的”Gemini Nano v3 模型上。基于之前的方法，如EAGLE框架 (https://arxiv.org/pdf/2401.15077)和确信自适应语言建模 (Confident Adaptive Language Modeling, CALM) (https://research.google/blog/accelerating-text-generation-with-confident-adaptive-language-modeling-calm/)，我们设计了新的架构组件，专门针对移动环境最大化这些效率提升。我们最近的公告强调了使用 MTP 加速Gemma 4 (https://blog.google/innovation-and-ai/technology/developers-tools/multi-token-prediction-gemma-4/)，并将其提供给开发者。 今天的文章探讨了边缘计算独特且极端的限制条件。这种方法最近已部署到 Pixel 9 和 10 系列上，作为一种开箱即用的加速手段。对用户而言，这意味着 AI 通知摘要和校对等功能生成文本的速度显著提升，同时能耗更低。对开发者而言，它消除了一个主要的痛点：无需为每个新任务分别微调那些占用大量内存的独立草稿模型，即可实现高速的端侧 AI。 ## “晚退出”策略 MTP 建立在推测解码 (speculative decoding) (https://research.google/blog/looking-back-at-speculative-decoding/) 的演进之上。在传统设置中，生成*N*个 token 需要对大模型进行*N*次前向传递。推测解码将这一过程分解为两步： 1. *草稿阶段*：一个更小、更快的近似模型（“草稿器”）生成一个较短的候选 token 序列（例如 3 个 token）。 2. *验证阶段*：一个大模型（“验证器”）并行处理这些候选 token。如果候选 token 与大模型本应预测的 token 匹配，则被接受；如果不匹配，系统会回滚到第一个分歧点。 然而，这会导致一些低效问题。运行一个独立的“草稿器”模型（例如 128M 参数）会争夺有限的内存。此外，独立的草稿器对主模型的丰富内部状态是“盲视”的，它仅基于文本历史预测下一个 token，而没有利用主模型已经计算出的语义上下文。MTP 通过从独立架构转向集成架构来解决这些低效问题。它不是训练一个单独的小语言模型来起草 token，而是在主模型的最终层上附加一个轻量级的 Transformer 头，即 MTP 头。 这种架构使用了一个深度退出层进行草稿生成，利用了主模型骨干已经完成的工作。MTP 头获取主模型最终的高维激活（隐藏状态），并利用它们来自回归地预测一系列未来 token。 ## 冻结骨干的优势 虽然 MTP 头通常与主干网络协同预训练 (https://arxiv.org/pdf/2404.19737)——例如我们最近发布的 Gemma 4 模型——但在利用已部署的端侧基础模型时，这种做法是不可行的。相反，我们的工作重点是将草稿器头改造为独立于预训练管道运行。 我们采用一个完全训练好的 Gemini Nano v3 模型，冻结其权重，并在其最终层上附加一个密集的 Transformer 堆栈——MTP 头。我们仅训练这些参数，以最小化未来 token 的预测误差。通过冻结骨干网络，MTP 成为一种纯粹的效率优化手段，确保基础模型的能力或安全对齐不会退化。 由于在验证过程中丢弃了错误的草稿，最终输出与主模型逐位一致，从而允许我们在完全向后兼容的情况下推出效率更新。 ## 零拷贝架构 虽然标准 MTP 实现 (https://arxiv.org/pdf/2404.19737)通过在主模型和草稿器之间共享静态参数（如嵌入权重）来优化训练效率，但端侧推理面临更严格的瓶颈：动态内存。即使共享权重，如果草稿器独立处理上下文，它也会通过生成和维护自身的键值 (key-value) (https://en.wikipedia.org/wiki/Key%E2%80%93value_database) \(KV\) 缓存而承担“双重内存负担”。考虑到移动设备内存有限，避免这种冗余至关重要。 为了解决这个问题，我们设计了一种零拷贝架构，其中 MTP 头有效地利用主模型的状态。MTP 头不维护自己的历史，而是被设计为直接交叉关注主模型的冻结 KV 缓存。这使得草稿器能够查询主干网络已经计算出的“记忆”和上下文，而无需重复。 这种设计带来了两个效率提升。首先，它消除了草稿器的预填充延迟：通过利用现有的缓存，MTP 头不需要额外的时间来处理提示。其次，它减少了运行时内存占用。与独立的草稿器相比，通过节省草稿器嵌入查找表、预填充点注意力变体以及应用特定的调优参数，我们观察到每个实例节省了 130MB 内存。 ## 释放更丰富的表征 在我们的实验中，我们发现 MTP 草稿器始终生成更准确的 token 预测，与参数数量相当的“独立草稿器”相比，在 Pixel 9 设备上实现了 50% 或更高的加速\[aef552\]，具体取决于任务。 这种性能差距源于 MTP 能够访问更丰富的表征。与将主模型视为黑盒的独立草稿器不同，MTP 头直接利用已经由更大骨干网络处理过的最终激活： - *指令遵循*：在具有复杂约束的摘要或重写等任务中，MTP 显著优于独立微调的草稿器。 - *可预测的文本结构*：对于具有高结构可预测性的任务（例如智能回复），MTP 头有效学习了主模型的句法模式，实现了 token 接受率高达 55% 的提升。 ## 实际影响 为了在 Pixel 9 和 10 设备上部署 MTP，我们重新设计了端侧推理栈，以处理验证阶段和草稿阶段之间的复杂依赖关系。 结果验证了架构选择的正确性。在生产工作负载中，例如 AI 通知摘要和校对，MTP 平均每次推理传递能正确预测近两个额外的 token。此外，更少的验证步骤意味着更少的时间需要唤醒重处理器，从而降低能耗并改善电池续航。 ### 未来方向 我们期待将 MTP 集成到未来的 Pixel 设备上，并探索替代架构——包括并行解码和无辅助头范式——以进一步降低草稿延迟，并在严格的移动约束下增加同时验证的 token 数量。 我们也在研究如何更高效地处理语言生成中固有的歧义性。虽然标准推测解码假设单一的最佳未来路径，但我们正在开发允许模型并行探索分支可能性的技术。这旨在最大化即使在不确定上下文中接受长序列的概率。此外，我们正在研究验证宽松度：针对特定用例放宽草稿与验证之间的严格精确 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 团队提供的审阅、宝贵反馈和支持。*