In Gemma 4 ist die Multi-Token Prediction (MTP) die spezifische Architektur, die für eine hocheffiziente spekulative Decodierung verwendet wird. Die spekulative Decodierung ist eine Technik, um die Inferenz in großen Sprachmodellen zu beschleunigen. Anstatt sich ausschließlich auf das große Zielmodell zu verlassen, um Tokens autoregressiv zu generieren (ein Token nach dem anderen, wobei jedes neue Token von den vorherigen abhängt), sagt ein kleineres, schnelleres „Entwurfsmodell“ mehrere Tokens voraus. Das Zielmodell überprüft diese Entwurfstokens dann parallel. Wenn das Zielmodell ein Entwurfstoken ablehnt, wird trotzdem das richtige Token für diese Position erstellt (damit dieser Schritt nicht verschwendet wird). Das Entwurfsmodell setzt die Vorhersage mit diesem neuen richtigen Token fort.
In Gemma 4 wird MTP implementiert, indem das Basismodell um dieses kleinere, schnellere Entwurfsmodell erweitert wird. Dieses Entwurfsmodell ist nicht unabhängig, da es die Eingabe-Embedding-Tabelle mit dem Zielmodell teilt und direkt auf den Aktivierungen der letzten Ebene aufbaut. Dies führt zu einer erheblichen Beschleunigung der Decodierung und garantiert gleichzeitig die gleiche Qualität wie die standardmäßige autoregressive Generierung. Daher eignen sich diese Prüfpunkte perfekt für Anwendungen mit geringer Latenz und auf Geräten.
Bei der spekulativen Decodierung werden mehrere Tokens entworfen und in einem einzigen Vorwärtsdurchlauf überprüft. Bei dichten Modellen werden für jedes Token dieselben Gewichte verwendet. Die Überprüfung mehrerer Entwurfstokens verursacht daher nur einen minimalen Mehraufwand. Mixture of Experts-Modelle (MoE) wie Gemma 4 26B A4B funktionieren anders. Jedes Token kann unterschiedliche Experten aktivieren. Daher kann es erforderlich sein, zusätzliche Expertengewichte aus dem Speicher zu laden, um die Entwurfstokens zu überprüfen. Dadurch werden die Vorteile des Entwurfs wieder zunichte gemacht. Bei größeren Batchgrößen gibt es in der Regel mehr Überschneidungen bei den aktivierten Experten in den Sequenzen, wodurch die Wiederverwendung geladener Gewichte verbessert wird. Bei einer Batchgröße von 1 ist diese Überschneidung begrenzt. Daher kann der 26B A4B-Entwurf auf Hardwareplattformen ohne gute Parallelität keine Geschwindigkeitssteigerung erzielen.
MTP-Verbesserungen
In Gemma 4 wurden mehrere Verbesserungen an der standardmäßigen spekulativen Decodierungspipeline eingeführt, um die Qualität der Entwurfstokens und die Effizienz zu verbessern:
- Gemeinsame Eingabe-Embeddings: Das Entwurfsmodell teilt die Eingabe-Embedding Tabelle mit dem Zielmodell.
- Zielaktivierungen: Das Entwurfsmodell verwendet die Aktivierungen aus der letzten Ebene des Zielmodells, verkettet sie mit den Token- Embeddings und projiziert sie auf die Dimension des Entwurfsmodells.
- Effizienter Embedder: Um die aufwendige Vorhersage über das gesamte Vokabular zu vermeiden, gruppiert das Modell ähnliche Tokens in Clustern. Zuerst werden die wahrscheinlichsten Cluster identifiziert und dann werden die endgültigen Berechnungen nur auf die Tokens in diesen ausgewählten Clustern beschränkt (nur E2B und E4B).