techlist.io
KREN
google

Accelerating Gemini Nano models on Pixel with frozen Multi-Token Prediction (opens in new tab)

k8stransformeredge-computingspeculative-decodinggemini-nanomulti-token-predictionon-device-inference

Google introduces a way to add Multi-Token Prediction (MTP) to already-deployed, frozen Gemini Nano models, accelerating on-device generation without a separate drafting model. The approach attaches a lightweight Transformer head to the existing model, reuses its hidden states and KV cache, and preserves identical final outputs through verification. On Pixel 9 and 10 devices, it delivers faster generation, lower memory use, and reduced energy consumption for features such as Notification Summaries and Proofread.

The Mobile Inference Bottleneck

  • Autoregressive models generate one token at a time, creating latency and underusing mobile hardware.
  • Phones face strict RAM and energy constraints that make conventional acceleration techniques difficult.
  • A standalone speculative-decoding drafter consumes additional memory and must independently process the prompt.

A “Late Exit” MTP Strategy

  • Speculative decoding uses:
    • A small drafter to propose several tokens.
    • The large model to verify those tokens in parallel.
  • MTP replaces the separate drafter with a lightweight Transformer head attached near the end of the main model.
  • The head uses the backbone’s high-dimensional activations to predict future tokens, benefiting from semantic context already computed by the larger model.

Retrofitting a Frozen Backbone

  • Google freezes the fully trained Gemini Nano v3 weights and trains only the attached MTP head.
  • This avoids retraining or fine-tuning the production foundation model.
  • Incorrect draft tokens are discarded during verification, so the final output remains bit-for-bit identical to the original model.
  • The method therefore improves efficiency without changing the model’s capabilities or safety alignment.

Zero-Copy Memory Architecture

  • The MTP head cross-attends directly to the backbone’s existing KV cache instead of maintaining a duplicate cache.
  • This eliminates separate prompt-prefill work for the drafter.
  • It also removes redundant embedding tables, attention variants, and application-specific tuning parameters.
  • Compared with a standalone drafter, the design saves up to 130 MB per instance.

Accuracy and Speed Improvements

  • Access to the backbone’s richer internal representations makes MTP predictions more accurate than those of similarly sized standalone drafters.
  • Instruction-following tasks such as summarization and rewriting show especially strong gains.
  • For predictable formats such as smart replies, token acceptance improved by up to 55%.
  • Pixel 9 experiments showed speedups of 50% or more depending on the task.

Production Impact

  • The updated inference stack coordinates drafting and verification on Pixel 9 and 10 devices.
  • In workloads including AI Notification Summaries and Proofread, MTP predicts nearly two additional tokens per inference pass on average.
  • Fewer verification cycles reduce processor wake-ups, improving latency and battery efficiency.

MTP provides a practical way to accelerate existing on-device models without maintaining a separate drafter for every application. Reusing the frozen backbone’s computations and memory state makes it particularly well suited to mobile deployments where RAM, energy, and backward compatibility are critical.