Speculative Decoding

2 posts

google3 min readCurated summary

Accelerating Gemini Nano models on Pixel with frozen Multi-Token Prediction

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.

Read original(opens in new tab)
googleOriginal article

Speculative cascades — A hybrid approach for smarter, faster LLM inference (opens in new tab)

Speculative cascades represent a hybrid inference method that integrates the cost-efficiency of model cascades with the latency-reducing benefits of speculative decoding. By utilizing a smaller drafter model to generate token sequences that are verified in parallel by a larger expert model, this approach allows for high-speed generation while maintaining flexible quality standards. The result is a system that achieves superior cost-quality trade-offs and higher speed-ups than either traditional cascading or standard speculative decoding alone. ### Limitations of Cascades and Speculative Decoding * **Sequential Bottlenecks in Cascades:** Traditional cascades use a deferral rule to decide if a small model can handle a prompt. If the small model is not confident, the system waits for it to finish before starting the large model from scratch, wasting significant time. * **Strict Matching in Speculative Decoding:** This method requires the large model to verify the small model’s tokens. Even if the small model produces a factually correct and high-quality response, the large model will reject the entire draft if the tokens do not match its own preferred output exactly. * **Trade-off Divergence:** Cascades prioritize reducing computational costs but suffer from latency when deferring, while speculative decoding prioritizes speed but often performs redundant work because it mandates identical output to the larger model. ### The Speculative Cascades Mechanism * **Parallel Verification with Deferral:** Speculative cascades use the parallel processing of speculative decoding but introduce a flexible decision rule. The system can choose to accept the smaller model’s draft even if it differs from the larger model’s prediction, provided it meets a confidence threshold. * **Flexible Token Matching:** Unlike standard speculative decoding, which often relies on strict token-by-token matching, speculative cascades allow for "probabilistic matches" or quality-based acceptance to prevent unnecessary rejections. * **Resource Optimization:** By strategically deferring to the smaller model for certain segments of the generation, the system reduces the total work required from the expensive expert model without losing the speed of parallel execution. ### Empirical Results and Performance * **Model Testing:** The approach was validated using Gemma and T5 models across diverse language tasks, including reasoning, coding, translation, and question answering. * **Superior Trade-offs:** Testing showed that speculative cascades consistently outperformed baselines in cost-quality metrics, providing faster inference without the strict "all-or-nothing" quality constraints of speculative decoding. * **Task Versatility:** The hybrid method proved effective across both creative tasks (like summarization) and factual tasks (like math or coding), where different levels of "correctness" are acceptable. Speculative cascades offer a practical path for scaling LLM deployments by balancing the high cost of large models with the need for low-latency user experiences. Developers looking to optimize inference should consider this hybrid approach to capture the efficiency of small models while retaining the oversight of larger, more capable ones.