Gemma 4 では、マルチトークン予測（MTP）は、非常に効率的な投機的デコードを可能にするために使用される特定のアーキテクチャです。投機的デコードは、大規模言語モデルの推論を高速化する手法です。 大規模なターゲット モデルのみに依存してトークンを自己回帰的に生成する（一度に 1 つのトークンを生成し、新しいトークンは前のトークンに依存する）のではなく、より小型で高速な「ドラフト モデル」が複数のトークンを予測します。ターゲット モデルは、これらのドラフト トークンを並行して検証します。ターゲット モデルがドラフト トークンを拒否した場合でも、その位置に正しいトークンが生成されます（このステップが無駄にならないようにします）。ドラフト モデルは、その新しい正しいトークンから予測を再開します。
Gemma 4 は、この小型で高速なドラフト モデルでベースモデルを拡張することで MTP を実装します。このドラフト モデルは、ターゲット モデルと入力エンベディング テーブルを共有し、最後のレイヤのアクティベーションを直接構築するため、独立していません。これにより、標準の自己回帰生成とまったく同じ品質を保証しながら、デコード速度が大幅に向上します。そのため、これらのチェックポイントは、低レイテンシのオンデバイス アプリケーションに最適です。
投機的デコードは、複数のトークンをドラフトし、1 回の順方向パスで検証することで機能します。高密度モデルの場合、すべてのトークンに同じ重みが使用されるため、複数のドラフト トークンを検証してもオーバーヘッドは最小限に抑えられます。Gemma 4 26B A4B などの混合エキスパート（MoE）モデルは動作が異なります。各トークンは異なるエキスパートをアクティブにする可能性があるため、ドラフト トークンを検証するには、メモリから追加のエキスパートの重みを読み込む必要があり、ドラフトによるメリットが相殺される可能性があります。バッチサイズが大きいほど、シーケンス間でアクティブ化されたエキスパートの重複が多くなり、読み込まれた重みの再利用が改善されます。バッチサイズが 1 の場合、この重複は制限されます。そのため、並列処理が適切でないハードウェア プラットフォームでは、26B A4B ドラフトが高速化されない可能性があります。
MTP の拡張機能
Gemma 4 では、ドラフト トークンの品質と効率を向上させるために、標準の投機的デコード パイプラインにいくつかの拡張機能が導入されています。
- 共有入力エンベディング: ドラフト モデルは、ターゲット モデルと入力エンベディング テーブルを共有します。
- ターゲット アクティベーション: ドラフト モデルは、ターゲット モデルの 最後のレイヤのアクティベーションを使用し、トークン エンベディングと連結して、ドラフト モデルのディメンションにダウン プロジェクトします。
- 効率的なエンベッダー: 語彙全体で予測するというコストのかかるオペレーションを回避するため、モデルは類似のトークンをクラスタにグループ化します。まず、最も可能性の高いクラスタを特定し、最終的な計算を、選択したクラスタ内のトークンのみに制限します（E2B と E4B のみ）。