Les modèles d’IA apprennent à diviser les tâches, réduisant ainsi les temps d’attente pour les invites complexes

À mesure que les grands modèles de langage (LLM) comme ChatGPT continuent de progresser, les attentes des utilisateurs à leur égard ne cessent de croître, notamment en ce qui concerne la rapidité avec laquelle ils peuvent répondre à nos invites de plus en plus complexes demandant des réponses à des problèmes et des tâches toujours plus difficiles.
Les LLM conventionnels s’appuient sur le concept de « décodage autorégressif », où chaque élément (« token ») d’une séquence est prédit à partir des sorties générées précédemment. Cette approche entraîne inévitablement des retards pour les requêtes plus complexes, bien que les chercheurs aient tenté d’atténuer ce problème grâce à des projets exploitant plus efficacement le parallélisme des puces informatiques multicœurs.
Par exemple, le décodage spéculatif utilise un modèle de brouillon rapide pour proposer des tokens qui sont ensuite vérifiés en parallèle par un modèle plus lent et de meilleure qualité. Une nouvelle classe de méthodes exploite plutôt « l’indépendance sémantique », identifiant des schémas syntaxiques tels que des puces et les développant chacun en parallèle. Cependant, elles s’appuient sur des heuristiques syntaxiques élaborées à la main, fragiles et souvent défaillantes lorsque les réponses s’écartent du format attendu.
Ces lacunes ont incité les chercheurs du Laboratoire d’informatique et d’intelligence artificielle (CSAIL) du MIT et de Google à utiliser une approche de décodage parallèle basée sur l’apprentissage. Au lieu de s’appuyer sur des règles fixes, leur méthode entraîne les LLM à reconnaître l’indépendance sémantique, c’est-à-dire à identifier et décoder en parallèle des fragments de texte sémantiquement indépendants.
Le résultat : PASTA
Plus précisément, l’annotation de structures parallèles (PASTA) de l’équipe CSAIL permet aux étudiants de LLM de générer du texte en parallèle, accélérant ainsi considérablement leurs temps de réponse. Contrairement aux tentatives précédentes qui s’appuyaient sur des règles rigides et codées manuellement pour identifier des segments de texte indépendants, PASTA apprend aux étudiants de LLM à comprendre et à exprimer intrinsèquement ces opportunités de parallélisation dans leurs propres réponses. Cette approche, appelée décodage asynchrone appris, marque une évolution vers des modèles d’enseignement permettant d’orchestrer sa propre stratégie de décodage parallèle.
« Les LLM traditionnels sont comme un cuisinier qui prépare des lasagnes étape par étape », explique Tian Jin, auteur principal d’un nouvel article sur le projet qui sera présenté en juillet à la Conférence internationale sur l’apprentissage automatique (ICML) à Vancouver. « PASTA apprend au cuisinier à reconnaître quand différentes parties des lasagnes peuvent être préparées simultanément, par exemple en mélangeant un sous-ensemble d’ingrédients pendant que le four préchauffe, ce qui accélère considérablement le processus. »
Cette innovation s’attaque à un goulot d’étranglement fondamental dans l’inférence LLM : la nature séquentielle du décodage entraîne souvent une sous-utilisation du matériel et des temps d’attente longs pour les utilisateurs. Les LLM actuels peuvent mettre quelques secondes, voire quelques minutes, à répondre aux requêtes des utilisateurs, un problème de latence que PASTA vise à résoudre.
Au cœur de PASTA se trouvent deux composants principaux : PASTA-LANG, un langage d’annotation qui permet aux LLM d’étiqueter des parties sémantiquement indépendantes de leurs réponses, et un interpréteur qui exploite ces étiquettes pour orchestrer le décodage parallèle lors de l’inférence. Comme l’explique Tian Jin, PASTA-LANG peut être considéré comme un ensemble d’instructions que le LLM écrit pour lui-même, marquant les sections de sa sortie pouvant être traitées simultanément. L’interpréteur lit ensuite ces instructions et gère la génération parallèle de ces sections.
L’équipe a entraîné les LLM à générer ces annotations PASTA-LANG grâce à un processus de réglage fin en deux étapes. Cet entraînement optimise non seulement la vitesse de décodage, mais maintient également, voire améliore, la qualité des réponses générées. Cette double optimisation constitue une avancée significative, car elle permet des améliorations continues de la vitesse et de la qualité à mesure que les capacités de calcul d’entraînement augmentent.
Lors d’expériences menées avec PASTA sur le benchmark AlpacaEval utilisé, le modèle auto-parallélisant de l’équipe a montré des accélérations moyennes géométriques atteignant près de 2x, avec des variations mineures de la qualité de réponse (d’un gain de 2 % à une baisse de 7 %). Cela signifie que les utilisateurs peuvent s’attendre à des réponses presque deux fois plus rapides sans perte notable de précision ou de cohérence.
« C’était surprenant de voir un LLM orchestrer lui-même son propre comportement au moment de l’inférence », explique Tian Jin. « C’était éclairant – et d’une certaine manière magique – de voir comment l’augmentation de la puissance de calcul de ces algorithmes produit un comportement d’auto-orchestration de plus en plus sophistiqué. »
Cette recherche met en évidence un défi majeur dans ce domaine : trouver le juste équilibre entre vitesse et qualité. Des méthodes antérieures, telles que le Skeleton-of-Thought (SoT) et l’APAR, tentaient un décodage parallèle en recherchant des structures syntaxiques spécifiées manuellement, comme des puces ou des paragraphes. Cependant, ces méthodes étaient souvent rigides et imprécises, ne parvenant pas à identifier les opportunités de parallélisation lorsque les réponses s’écartaient, même légèrement, des modèles attendus. L’approche basée sur l’apprentissage de PASTA, en revanche, offre une solution plus robuste et évolutive.
« Il s’agit de permettre au LLM de générer plus intelligemment du contenu », précise Tian Jin, doctorant au CSAIL. « Au lieu d’essayer de deviner où il peut fonctionner en parallèle, nous apprenons au LLM à identifier lui-même ces opportunités, instantanément. »
L’équipe se montre optimiste quant aux implications plus larges de PASTA. La possibilité de réduire significativement la latence de décodage LLM pourrait réduire les besoins en ressources de calcul, rendant ces puissants modèles d’IA plus accessibles et abordables pour un plus large éventail d’utilisateurs et d’applications.
« Nous avons essentiellement conçu un protocole permettant à un LLM de s’auto-optimiser », assure Tian Jin. « En améliorant l’efficacité de l’inférence des LLM, PASTA pourrait réduire considérablement les besoins en ressources de calcul et améliorer l’accessibilité des LLM. »
Tian Jin a piloté le projet aux côtés de ses deux conseillers pédagogiques, les professeurs du MIT Michael Carbin et Jonathan Ragan-Kelley. Parmi les autres co-auteurs de l’article figurent Ellie Y. Cheng et Zack Ankner du CSAIL, ainsi que les chercheurs de Google Suvinay Subramanian, Nikunj Saunshi, Blake M. Elias et Amir Yazdanbakhsh.
Le travail a été soutenu en partie par la Fondation Sloan et SRC JUMP 2.0 (CoCoSys), ainsi que par l’équipe Google CoreML Performance, Google Research et Google DeepMind.
https://www.csail.mit.edu/news/speeding-llm-inference-parallelism
