Google Releases Gemma 4 Inference Acceleration Solution, Boosting Speed 3x
In May 2024, Google announced the open-sourcing of companion draft models for the Gemma 4 series. Designed specifically for accelerating inference, these models leverage a core technology called Speculative Decoding. By using a lightweight draft model in tandem with the main Gemma 4 model, they can achieve up to a threefold increase in inference speed while fully preserving the output quality and capabilities of the original Gemma 4 model.
Speculative Decoding: The Core Mechanism for Acceleration Without Sacrificing Accuracy
Speculative decoding is an advanced inference optimization technique. Its fundamental principle is to use a small “draft model,” which is much less computationally expensive than the main model, to pre-generate a sequence of candidate tokens. This draft sequence is then verified in a single batch by the more powerful and accurate “main model.”
The specific workflow is as follows:
- Draft Generation: The lightweight draft model performs Multi-Token Prediction (MTP) to generate a “draft” sequence containing multiple tokens at once.
- Batch Verification: The main Gemma 4 model receives the entire draft sequence and performs a single forward pass to verify its accuracy.
- Conditional Acceptance:
- If the main model fully accepts the draft content, all draft tokens are accepted at once, and the next token generated by the main model during this verification pass is appended. This allows multiple tokens to be generated in a single computation cycle.
- If the main model finds a mismatched token during verification, it accepts all tokens from the beginning of the sequence up to the first mismatch. The mismatched token is replaced by the main model’s prediction, and all subsequent draft tokens are discarded.
Since every token in the final output is either verified or generated by the main Gemma 4 model, this method ensures that the output quality is identical to using the main model alone, achieving “lossless acceleration.”
Key Optimization Techniques for the Draft Model
To enable the draft model to “guess” the main model’s output more quickly and accurately, Google has integrated three key technologies into its design.
Target Activations
To improve the draft model’s prediction accuracy, this technique allows it to “reference” the internal state of the main model. When generating the first token of the draft, the draft model concatenates the output embedding from the main model’s last layer with its own token embedding, and then maps the result to its own hidden layer dimension. This process effectively provides a strong “hint” from the main model to the draft model. For subsequent tokens in the draft sequence, the draft model concatenates the embedding vector it generated in the previous step.
KV Cache Sharing
To avoid redundant computations, the draft model is designed to directly utilize the Key-Value (KV) Cache already computed by the main model. Specifically, the local attention layers of the draft model reuse the most recent local KV cache from the main model, while its global attention layers connect directly to the global KV cache of the Gemma 4’s final layer. This sharing mechanism significantly reduces the computational overhead of the draft model and is a crucial part of the overall speedup.
Efficient Embedder
This technique primarily addresses the computation bottleneck of the Language Model (LM) Head in smaller models (like Gemma 4 E2B/E4B) deployed on edge devices. Traditional models require multiplying the hidden state with the entire vocabulary weight matrix, which is computationally intensive. The Efficient Embedder works by clustering all token embeddings, grouping semantically similar tokens into a cluster, and calculating a representative embedding for each cluster. During prediction, the model first predicts a cluster and then selects a token within that cluster, thereby decomposing a large matrix multiplication into smaller, more efficient computational steps. For larger models with more parameters (like Gemma 4 26B/31B), the Decoder part is the main bottleneck, so the gains from this technique are relatively limited.
