Consistency LLM : transformer les LLM en décodeurs parallèles pour accélérer l’inférence de 3,5×

Consistency Large Language Models (CLLM)

• Les LLM ont traditionnellement été considérés comme des décodeurs séquentiels fonctionnant en décodant les tokens un par un de manière successive
• Ce billet montre que des LLM préentraînés peuvent être facilement convertis en décodeurs parallèles efficaces
• Consistency Large Language Models (CLLM) constitue une nouvelle famille de décodeurs parallèles capable de décoder efficacement des séquences de n tokens à chaque étape d’inférence afin de réduire la latence
• CLLM imite le processus cognitif humain qui consiste à former une phrase complète dans sa tête avant de prononcer les mots un par un, et peut être entraîné efficacement simplement en effectuant un fine-tuning d’un LLM préentraîné
• CLLM est entraîné à effectuer un décodage parallèle en mappant une séquence de n tokens initialisée aléatoirement vers le même résultat que le décodage auto-régressif (autoregressive, AR) en aussi peu d’étapes que possible
• Selon les résultats expérimentaux, CLLM a montré une amélioration de la vitesse de génération de 2,4× à 3,4× par rapport aux décodeurs AR classiques, avec des performances comparables voire supérieures à des techniques d’inférence rapide comme Medusa2 ou Eagle
• CLLM peut atteindre ce gain de performance sans coût mémoire supplémentaire

Contexte et limites du décodage Jacobi

• Les LLM génèrent les tokens un par un avec le décodage AR, ce qui entraîne une latence élevée pour les réponses longues
• Le décodage Jacobi dérive des méthodes de résolution d’équations non linéaires de Jacobi et de Gauss-Seidel, et il a été démontré qu’il est équivalent à la génération AR avec greedy sampling
• Le décodage Jacobi reformule le processus de génération séquentielle en un système de n équations non linéaires à n variables fondé sur des itérations de Jacobi, ce qui permet un traitement parallèle
• À chaque étape d’itération, il est possible de prédire au moins un token correct, ce qui peut potentiellement accélérer le décodage AR
• Cependant, en pratique, les LLM entraînés en AR produisent rarement les bons tokens lorsque les tokens précédents contiennent des erreurs ; dans la plupart des itérations Jacobi, on n’obtient donc qu’une seule correction sur la séquence de n tokens, ce qui génère une trajectoire Jacobi plus longue
• Le décodage lookahead ou le speculative decoding tentent d’atténuer cette inefficacité du décodage Jacobi, mais ils entraînent un coût mémoire supplémentaire à l’inférence, contrairement à CLLM

Méthode d’entraînement de CLLM

• L’entraînement de CLLM se compose principalement de deux parties : la préparation des trajectoires Jacobi, puis l’optimisation des pertes de consistency et AR
• Lors de la préparation des trajectoires Jacobi, on effectue un décodage Jacobi séquentiel par blocs de n tokens jusqu’à la génération complète d’une séquence de réponse de l tokens, et la séquence générée dans chaque trajectoire est considérée comme un élément de données
• Pendant l’entraînement, on optimise conjointement une perte de consistency et une perte AR : la perte de consistency garantit la prédiction simultanée de plusieurs tokens, tandis que la perte AR empêche CLLM de trop s’éloigner du LLM cible afin de préserver la qualité de génération
• La perte de global consistency (GC) minimise la distance entre un point arbitraire de la trajectoire Jacobi et le point fixe, afin d’encourager CLLM à prédire le point fixe depuis n’importe quel point de la trajectoire
• La perte de local consistency (LC) pousse les états adjacents de la trajectoire Jacobi à produire la même sortie
• La perte AR inclut une perte auto-régressive classique fondée sur les résultats générés par le LLM cible ; son objectif est d’éviter que CLLM ne s’écarte de la distribution du LLM cible

Résultats expérimentaux

• Les expériences incluaient des tâches spécialisées dans des domaines ciblés comme Spider (text-to-SQL), Human-Eval (complétion de code Python) et GSM8k (mathématiques), ainsi que des défis de dialogue ouvert plus larges comme MT-bench
• CLLM a obtenu les plus grands gains de vitesse par rapport au modèle cible et a atteint des accélérations similaires ou meilleures que Medusa2 sans coût supplémentaire à l’inférence
• Sur MT-bench, CLLM a atteint quasiment le même gain de vitesse que lorsqu’il est combiné à Medusa2, tout en offrant une meilleure adaptabilité et une meilleure efficacité mémoire
• Le coût de fine-tuning de CLLM reste modéré et, même avec de grands jeux de données, l’utilisation d’environ 10 % du dataset pour générer les trajectoires Jacobi permettait déjà d’obtenir un gain de vitesse d’environ 2,5×
• CLLM peut prédire correctement plusieurs tokens consécutifs en une seule itération Jacobi grâce au phénomène de fast forwarding
• CLLM montre une capacité de tokens stationnaires, c’est-à-dire qu’il peut prédire à l’avance les bons tokens et les maintenir inchangés même lorsque les tokens précédents comportent des erreurs
• Grâce à l’apprentissage, CLLM acquiert des concepts linguistiques essentiels tels que les collocations ; il peut ainsi inférer la structure depuis n’importe quel point de la trajectoire Jacobi et prédire plusieurs mots simultanément afin de minimiser le nombre d’itérations

L’avis de GN⁺

• CLLM semble résoudre efficacement le problème de forte latence du décodage AR des LLM existants en s’appuyant sur le décodage Jacobi. Le fait qu’il obtienne un gain de vitesse via un décodage parallélisé sans coût mémoire supplémentaire est particulièrement impressionnant

• La méthode d’entraînement de CLLM, qui consiste à effectuer un fine-tuning d’un LLM existant via une perte de consistency, paraît relativement simple, mais elle semble significative dans la mesure où elle améliore fortement les performances du décodage parallèle en apprenant l’une des propriétés importantes du langage, à savoir les collocations

• Cela dit, comme CLLM suppose un greedy sampling, des recherches supplémentaires semblent nécessaires pour vérifier s’il fonctionne bien avec des stratégies de décodage plus variées. De plus, les expériences étant pour l’instant limitées à l’anglais, il faudra aussi vérifier sa capacité de généralisation à d’autres langues

• CLLM semble être une approche pratique pour accélérer le temps de réponse des LLM. Il pourrait bien s’appliquer à des tâches nécessitant une forte réactivité, comme la recherche web ou les chatbots

• À titre personnel, on peut espérer que la méthode d’apprentissage par consistency de CLLM puisse s’appliquer non seulement aux LLM comme GPT, mais aussi à d’autres modèles génératifs, comme les modèles de génération d’images ou de synthèse vocale. Il serait souhaitable que l’idée de CLLM contribue à l’avenir à améliorer l’efficacité de nombreux modèles génératifs

[GN⁺]

Commentaires sur Hacker News

• Dans les cours de dessin libre, s’exercer à dessiner de façon répétée avec une limite de temps améliore fortement le niveau. Le fait de s’entraîner à dessiner vite aide au contraire à mieux saisir les proportions et les contours avec précision.
• Le décodage de Jacobi est identique au décodage auto-régressif greedy, mais en pratique il faut régler la température d’échantillonnage au-dessus de 0 pour éviter les répétitions et des réponses trop génériques.
• Nous réaliserons bientôt qu’il n’est pas nécessaire d’entraîner les modèles. Il suffit d’un bon indexing et d’un bon échantillonnage. Les LLM sont essentiellement l’équivalent d’une base de données de datasets dotée d’une excellente interface NLP.
• Le gain de performances en inférence du papier CLLM est remarquable. Le coût du fine-tuning est raisonnable (environ 0,01 % du coût de préentraînement) et l’amélioration des performances est assez constante.
• Les trajectoires de Jacobi excluent les chemins à haute température, ce qui est positif pour la recherche de données mais peut être défavorable à la maximisation de la créativité.
• Ce serait bien d’avoir une fonction « demander à un expert IA » pour comprendre pourquoi un LLM ne répond pas de manière déterministe à un même prompt.
• Je me demande si l’accélération obtenue par Groq (https://groq.com/) vient de cette méthode ou d’une autre.