发生了什么

一位使用 Google 的 LiteRT API 开发 Android 应用的开发者发现,Gemma 4 包含未在公开发布模型中暴露的多 token 预测(MTP)权重。这一问题在使用 Google Pixel 9 测试设备加载模型时浮出水面,设备抛出了引用“mtp 权重具有不兼容的张量形状”的错误。该错误直接指向嵌入在 LiteRT 文件中的额外预测头部——这些头部专为推测解码设计,旨在加速 token 生成。

在将发现发布到 Hugging Face 上的 Gemma 4 模型讨论区(google/gemma-4-E4B-it)后,该开发者收到了一名 Google 员工的确认,证实 Gemma 4 确实在内部包含了 MTP 架构。官方给出的禁用理由是“确保兼容性和广泛的可用性”。目前尚未提供 MTP 支持何时或是否会被正式启用的时间表。

受影响的模型是 Gemma 4 E4B(40 亿有效参数的混合专家模型),这是 Google 的端侧变体。这一发现也重新引发了对更大规模 Gemma 124B 模型被遗漏的关注,该模型曾短暂出现在 Google 研究员 Jeff Dean 的一条推文中,随后被撤回。

技术深度解析

多 token 预测是一种推理时技术,模型利用辅助“草稿”头部同时预测多个未来 token,而不是每次前向传播生成一个 token。这与推测解码密切相关,后者由较小的草稿模型或内部草稿头部提出几个 token,主模型随后在单次前向传播中进行验证——在不改变输出质量的情况下显著提高吞吐量。

在 Gemma 4 的案例中,MTP 头部似乎作为额外的张量权重被烘焙到 LiteRT flatbuffer 文件中。LiteRT(前身为 TensorFlow Lite)使用静态计算图格式,这意味着图结构和权重形状在导出时即已固定。不兼容的张量形状错误表明运行时未配置为将激活路由通过这些头部,导致形状不匹配而非干净地跳过。

与 Medusa(添加多个独立训练的外部草稿头部)或 DeepSeek 的 MTP 实现(使用共享 Transformer 层馈入顺序预测头部)等方法不同,Gemma 4 的 MTP 架构细节尚未公开。理论上,可以使用 flatc 等工具逆向工程 LiteRT 计算图,反编译 flatbuffer 模式,或使用 Netron 可视化计算图并识别孤立的 MTP 子图。

作为对比,Meta 的 Llama 3 系列不包含原生 MTP 头部,而是依赖外部推测解码配合单独的草稿模型。Google 自家的 Gemini Nano 通过 LiteRT 采用了类似的端侧优化路径,这使得 Gemma 4 的隐藏 MTP 头部本应是同一管道的自然组成部分——如果它们被启用的话。

粗略的提取方法包括:

  • 使用 TFLite 模式和 flatc --json 反编译 .tflite flatbuffer
  • 识别与 MTP 头部权重名称关联的张量索引
  • 重构子图连接和输出形状
  • 修补模型,将最终隐藏状态路由至草稿头部

在没有访问训练检查点的情况下,这些被禁用的头部是完全训练的还是部分初始化的仍不得而知。

谁应该关注

针对 Android 硬件(特别是配备专用 NPU 能力的 Pixel 8 和 Pixel 9 设备)的端侧 AI 开发者与此最为相关。在像 Gemma 4 E4B 这样已经快速的 MoE 架构上启用 MTP,可以显著降低完全在端侧运行的实时应用(如键盘建议、语音助手和文档摘要)的延迟。

从事推测解码和高效推理的研究人员也应予以关注。如果 MTP 头部可以从 LiteRT 文件中恢复,这将提供罕见的视角,了解 Google 如何在生产级 MoE 模型中构建草稿头部——这些架构细节并未在任何公开的 Gemma 4 技术报告中涵盖。

通过 Hugging Face Transformers 或 llama.cpp 后端在服务器硬件上部署 Gemma 4 的 ML 工程师受到的直接影响较小,因为这些路径使用的是 SafeTensors 权重而非 LiteRT flatbuffers,且 MTP 头部似乎并不存在于那里。

本周行动指南

如果您想亲自调查这些隐藏的 MTP 权重,请从官方 Gemma 4 E4B 发布的 LiteRT 模型文件开始:

  • 下载模型:huggingface-cli download google/gemma-4-E4B-it
  • 安装 Netron (netron.app) 并打开 .tflite 文件,可视化计算图并发现未连接的 MTP 子图
  • 或者,安装 flatbuffers:pip install flatbuffers,并使用 TensorFlow 仓库中的 TFLite 模式反编译二进制文件
  • 搜索包含“mtp”、“draft”或“speculative”字符串的张量名称以定位相关权重
  • 关注 Hugging Face 讨论区 huggingface.co/google/gemma-4-E4B-it/discussions/5 以获取社区更新

在完全理解子图结构之前,请勿修改生产应用构建——在不知道其预期输入归一化的情况下将激活路由至未测试的头部可能会产生垃圾输出。