Comment Gemma 4 accélère la génération avec la prédiction multi-jetons

Les grands modèles de langage comme Gemma 4 génèrent leurs réponses jeton après jeton. Ce mode de calcul crée une latence sensible dès que le modèle grossit. Google s’appuie donc sur des drafters spécialisés pour prévoir plusieurs jetons à l’avance et réduire ce délai, sans dégrader la qualité finale des sorties.

---

Le goulot d’étranglement mémoire des modèles de langage

L’inférence des grands modèles reste surtout limitée par la bande passante mémoire, plus que par la puissance de calcul brute. À chaque nouveau jeton, le processeur doit relire presque tous les paramètres depuis la mémoire vive. Les unités de calcul attendent, ce qui ralentit la génération.

Pourquoi le modèle semble-t-il bégayer lors de la génération

Un modèle de 31 milliards de paramètres occupe environ 62 Go en FP16. À chaque étape, le GPU calcule quelques millisecondes, puis attend les données. C’est ce va-et-vient qui donne l’impression d’un modèle qui « bégaie ».

Le mécanisme classique de prédiction jeton par jeton

La méthode traditionnelle, appelée Next-Token Prediction, calcule la probabilité du jeton suivant à partir du contexte, puis recommence. Chaque étape dépend de la précédente. Le traitement reste donc strictement séquentiel.

Principe du décodage spéculatif par multi-token prediction

Le décodage spéculatif sépare la proposition des jetons de leur validation. Un petit modèle, appelé drafter, anticipe plusieurs jetons pendant que le modèle principal les vérifie en une seule passe.

Comment un drafter MTP propose plusieurs jetons simultanément

Le drafter, entraîné pour cette tâche, propose trois ou quatre jetons à partir du contexte courant. Il le fait en quelques millisecondes, car il reste très léger. Le modèle cible examine ensuite ces propositions ensemble.

Le cycle de validation et de correction automatique

Si le modèle cible accepte les jetons proposés, plusieurs sont validés d’un seul coup. En cas de rejet partiel, il corrige le premier jeton erroné et relance le cycle. Le résultat final reste identique à celui d’une génération classique.

L’architecture native des drafters dans Gemma 4

Contrairement aux méthodes externes, les drafters de Gemma 4 sont intégrés dès la conception au modèle principal. Cette architecture facilite le partage des ressources de calcul.

Une structure légère de quatre couches seulement

Chaque drafter ne compte que quatre couches de transformeur. Son empreinte mémoire se limite à quelques centaines de mégaoctets face à un modèle cible de 31 milliards de paramètres. Cette légèreté permet de l’exécuter en parallèle sans surcharger le matériel.

Partage du cache de clés et valeurs avec le modèle cible

Le drafter réutilise directement le KV Cache déjà calculé par le modèle principal. Il n’a donc pas à recalculer les représentations des jetons précédents. Ce partage réduit les doublons en mémoire et en calcul.

Exploitation des activations de couche intermédiaire

Les activations produites par la dernière couche du modèle cible sont concaténées puis projetées pour affiner les prédictions du drafter. Ce lien direct aligne le petit modèle sur les représentations internes du grand modèle. Il améliore aussi le taux d’acceptation des jetons spéculés.

Gains de vitesse mesurés sur différents matériels

Google annonce des accélérations pouvant atteindre un facteur trois sur Gemma 4. Les résultats varient selon le type de modèle et le matériel utilisé.

Performances sur les modèles denses de 31 milliards

Sur le modèle dense 31B, la réutilisation des poids reste efficace et le gain approche souvent 3x en débit. Le drafter fait accepter en moyenne plus de deux jetons par tentative, ce qui réduit le nombre de passes sur le modèle principal.

Comportement particulier des modèles Mixture of Experts

Le modèle 26B A4B, de type Mixture of Experts, affiche des gains plus variables. Avec une taille de lot de un, chaque jeton peut activer des experts différents, ce qui limite le partage du cache. Sur des charges plus parallèles ou sur TPU, le débit triple toutefois.

Résultats concrets sur GPU, TPU et puces Apple

Sur les puces Apple Silicon M5 Max, les outils MTPLX font passer la vitesse de 28 à 63 jetons par seconde, soit un facteur 2,25. Sur les GPU NVIDIA récents et sur TPU v5, les mesures internes de Google confirment un débit triplé lorsque la taille de lot exploite le parallélisme des propositions spéculatives.

Mise en œuvre pratique et limites identifiées

Les drafters MTP sont publiés sous licence Apache 2.0 et peuvent être activés dans la plupart des environnements d’inférence standards. Selon le type de contenu généré, quelques précautions restent nécessaires.

Intégration dans vLLM, SGLang et Hugging Face

Dans vLLM, il suffit d’indiquer le modèle assistant avec le paramètre speculative_model. Avec Hugging Face Transformers, on charge le modèle cible puis le drafter correspondant en deux lignes. SGLang propose aussi une option native pour activer le décodage spéculatif sans toucher au code applicatif.

Optimisation pour les appareils mobiles et de bord

Sur Android et iOS, la bibliothèque LiteRT-LM, intégrée à Google AI Edge Gallery, permet d’exécuter les couples modèle-drafter avec une latence inférieure à la centaine de millisecondes par réponse. Un ajustement dynamique réduit le nombre de jetons spéculés dès que le taux d’acceptation passe sous le seuil défini.

Recommandations selon la nature du texte généré

La génération de code reste plus difficile à anticiper que le langage naturel. Google recommande de limiter les propositions à trois ou quatre jetons pour les tâches de programmation afin de limiter les rejets. Pour le texte conversationnel ou la rédaction, des propositions plus longues restent efficaces et rentabilisent mieux l’approche.