LLM : le décodage spéculatif triple la vitesse de génération

Le processus standard de génération de texte par les LLM

Dans le cadre de la génération de texte par un modèle de langage de grande taille (LLM), deux étapes essentielles sont généralement suivies. Tout d'abord, il y a le passage avant. Durant cette phase, le texte d'entrée est converti en tokens et introduit dans le modèle. Chaque couche du modèle transforme ces tokens, et à la fin, une distribution de probabilité est produite pour les tokens suivants possibles. Cela signifie que chaque token potentiel est associé à une probabilité.

Ensuite, vient la phase de décodage. Dans cette étape, le prochain token est sélectionné à partir de la distribution de probabilité. Ce choix peut se faire soit en optant pour le token avec la probabilité la plus élevée, un processus connu sous le nom de décodage glouton, soit en échantillonnant parmi les tokens les plus probables, une méthode appelée échantillonnage top-p ou nucleus.

Après la sélection d'un token, celui-ci est ajouté à la séquence d'entrée, et un nouveau passage avant est effectué pour générer le token suivant. Pour un modèle de grande taille, par exemple avec 70 milliards de paramètres, chaque token nécessite un passage complet à travers le modèle, ce qui rend le processus de génération long et coûteux en termes de temps.

Les modèles autoregressifs fonctionnent de manière séquentielle, comme des dominos. Ainsi, le token 100 ne peut être généré que lorsque tous les tokens précédents ont été traités. Chaque token exige un passage complet à travers le réseau, et générer 100 tokens à une vitesse de 20 ms par token entraîne un délai d'environ 2 secondes, chaque token devant attendre le traitement des précédents. Cela se traduit par une latence significative.

L'innovation du décodage spéculatif

Pour pallier ces limitations, le décodage spéculatif utilise deux modèles : un grand LLM, appelé modèle cible, et un modèle plus petit, souvent une version distillée, désigné comme modèle de brouillon. L'idée principale est que le modèle de brouillon génère rapidement des tokens qui sont plus simples et prévisibles, tandis que le modèle cible vérifie la précision, notamment pour les tokens plus complexes ou spécifiques à un domaine.

Le modèle de brouillon imite le comportement du modèle cible pour la majorité des tokens, mais le modèle cible agit comme un contrôleur de qualité pour garantir la précision globale de la sortie.

La stratégie du décodage spéculatif repose sur trois étapes clés :

• Brouillon : Générer rapidement K tokens avec le modèle de brouillon.
• Vérification : Effectuer un passage avant unique du modèle cible sur tous les K tokens en parallèle.
• Acceptation/Rejet : Accepter les tokens corrects et remplacer les incorrects par l'échantillonnage par rejet.

Cette méthode a été proposée par Google Research et Google DeepMind dans leur article "Accelerating LLM Decoding with Speculative Decoding".

En général, un modèle génère un token par passage avant. Cependant, il est possible d'alimenter plusieurs tokens dans un LLM et de les évaluer en parallèle lors d'un seul passage avant. Vérifier une séquence de tokens est à peu près aussi coûteux que de générer un seul token, tout en produisant une distribution de probabilité pour chaque token de la séquence.

Le principe d'acceptation et de rejet

L'étape suivante est l'échantillonnage par rejet, où il est décidé quels tokens seront conservés et lesquels seront rejetés. Chaque token est examiné individuellement, en comparant les probabilités p(x) et q(x) assignées respectivement par les modèles de brouillon et cible.

L'acceptation ou le rejet se fait selon une règle simple :

• Cas 1 : si q(x) >= p(x), le token est accepté.
• Cas 2 : sinon, le token est rejeté.

Scénario pessimiste

Dans le pire des cas, le premier token est rejeté, ce qui implique que tous les tokens suivants, bien que potentiellement pertinents, sont également rejetés. Chaque token étant dépendant de son prédécesseur, le modèle cible doit corriger le token x1, puis le modèle de brouillon génère un nouvel ensemble de 5 tokens, qui seront à nouveau vérifiés par le modèle cible.

Scénario optimiste

Dans le meilleur des cas, tous les tokens générés par le modèle de brouillon sont acceptés par le modèle cible. De plus, le modèle cible génère un token supplémentaire après le x5. Ainsi, dans ce scénario, nous obtenons K+1 tokens en une seule fois, ce qui représente un gain de vitesse significatif.

Le décodage spéculatif permet un gain de vitesse d'environ 2 à 3 fois en générant et vérifiant les tokens en parallèle. L'échantillonnage par rejet est crucial pour garantir que la qualité de la sortie reste conforme au modèle cible, malgré l'utilisation de tokens de brouillon.

Efficacité du passage

• Pire cas : Le premier token est rejeté, ce qui signifie qu'un seul token du modèle cible est accepté.
• Meilleur cas : Tous les tokens de brouillon sont acceptés, ce qui donne (tokens de brouillon) + (token du modèle cible) [K+1].

Selon l'article de DeepMind, il est conseillé de maintenir K = 3 et 4, ce qui a souvent permis d'obtenir un gain de vitesse de 2 à 2,5 fois par rapport au décodage autoregressif. Dans l'article de Google, le chiffre 3 a été recommandé, permettant d'atteindre un gain de vitesse de 2 à 3,4 fois.