Running large language model capabilities directly on phones — for tasks like instant notification summaries or local proofreading — avoids sending private data off-device but requires very efficient execution. Mobile devices have strict energy budgets and limited RAM, and standard autoregressive generation (one token per forward pass) can create processing and memory bottlenecks that slow responses and drain battery.
To reduce this bottleneck, Google has retrofitted Multi-Token Prediction (MTP) onto existing, frozen Gemini Nano v3 models. Rather than training a separate small drafter model, the technique appends a lightweight Transformer head (the MTP head) to the final layers of the main model so the drafter can autoregressively predict multiple future tokens using the backbone’s already-computed hidden states.
"Late exit" strategy and speculative decoding
MTP builds on speculative decoding approaches: conventionally, a smaller drafter proposes short candidate token sequences and a larger verifier checks them in parallel. Separate drafter models, however, compete for RAM and lack access to the main model’s internal state, basing predictions only on text history. MTP moves drafting into a deep exit layer of the main model, leveraging the backbone’s rich activations and avoiding a standalone drafter’s blindness to internal context.
Advantage of a frozen backbone
Instead of pretraining drafter heads together with the backbone — an approach used for some Gemma 4 releases — the retrofit focuses on already-deployed on-device foundation models. The process takes a fully trained Gemini Nano v3, freezes its weights, attaches a dense transformer stack as the MTP head, and trains only those extra parameters to minimize future-token prediction error. Because the backbone remains frozen, MTP acts strictly as an efficiency optimization without degrading base capabilities or safety alignment. Incorrect drafts are discarded during verification, so final outputs are bit-for-bit identical to the main model, preserving backward compatibility.
Zero-copy architecture to avoid duplicate memory use
On-device inference is constrained by dynamic memory. Even with shared static parameters, a drafter that maintains its own key-value (KV) cache duplicates memory use. The implemented zero-copy architecture lets the MTP head cross-attend directly to the main model’s frozen KV cache rather than maintaining a separate history. This removes drafter prefill latency (the head can use the existing cache immediately) and lowers runtime memory. The team reports about 130 MB saved per instance versus a standalone drafter by avoiding drafter embedding lookup tables, prefill attention variants, and some tuning parameters.
Richer representations yield better drafts and speedups
Experiments showed MTP drafters consistently produce more accurate token predictions, translating to speedups on Pixel 9 devices of 50% or more versus standalone drafters of comparable parameter counts. MTP’s performance improvement comes from access to richer representations already present in the backbone’s final activations:
- Instruction following: for summarization or constrained rewriting, MTP significantly outperformed standalone fine-tuned drafters.
- Predictable text structures: for tasks with high structural predictability (e.g., smart replies), the MTP head learned the main model’s syntactic patterns and achieved up to a 55% improvement in token acceptance.
In production workloads on Pixel 9 and 10 devices (such as AI Notification Summaries and Proofread), MTP correctly predicts on average nearly two additional tokens per inference pass. Fewer verification steps also reduce the time heavy processors need to run, lowering energy use and improving battery life.
Future directions
The team plans to integrate MTP on future Pixel devices and to explore alternative architectures — including parallel decoding and approaches without auxiliary heads — to further reduce draft latency and increase simultaneous token verification under mobile constraints. They are also investigating methods to handle generation ambiguity more efficiently, such as exploring branching possibilities in parallel rather than assuming a single best future path, and studying verification leniency (relaxing exact token match) for certain use cases to gain additional edge efficiencies.
Acknowledgements
This optimization effort lists contributors Filippo Galgani, Omri Homburger, Pooja Consul, Matthew Markwell, and Vivek Kumar. Some elements build on developments from the Gemini team in Google DeepMind: Tal Schuster, Ziwei Ji, Ivan Korotkov, and Ganesh Jawahar. Reviews and support came from Nadav Bar, Utku Evci, Nir Shabat, Joe Zou, and teams in Google Research, Google DeepMind, and Platforms & Devices.
