Lorsqu'une conversation homme-IA implique de nombreux cycles de dialogue continu, les puissants modèles d'apprentissage automatique à grande échelle qui pilotent les chatbots comme ChatGPT commencent parfois à s'effondrer, entraînant une détérioration rapide des performances des robots.
Une équipe de chercheurs du MIT et d'ailleurs a identifié une cause surprenante de ce problème et développé une solution simple qui permet à un chatbot de maintenir une conversation non-stop sans planter ni ralentir.
Leur méthode implique une modification du cache clé-valeur (qui est comme une mémoire de conversation) au cœur de nombreux grands modèles de langage. Dans certaines méthodes, lorsque ce cache doit contenir plus d'informations que ce dont il a la capacité, les premiers éléments de données sont supprimés. Cela peut entraîner l'échec du modèle.
En garantissant que ces premiers points de données restent en mémoire, la méthode des chercheurs permet à un chatbot de continuer à discuter quelle que soit la durée de la conversation.
La méthode, appelée StreamingLLM, permet à un modèle de rester efficace même lorsqu'une conversation s'étend sur plus de 4 millions de mots. Comparé à une autre méthode qui évite les plantages en recalculant constamment une partie des conversations passées, StreamingLLM a fonctionné plus de 22 fois plus rapidement.
Cela pourrait permettre à un chatbot de mener de longues conversations tout au long de la journée de travail sans avoir besoin d'être continuellement redémarré, ce qui permettrait à des assistants IA efficaces pour des tâches telles que la rédaction, l'édition ou la génération de code.
« Désormais, grâce à cette méthode, nous pouvons déployer de manière persistante ces grands modèles de langage. En créant un chatbot avec lequel nous pouvons toujours discuter et qui peut toujours nous répondre en fonction de nos conversations récentes, nous pourrions utiliser ces chatbots dans de nouvelles applications », explique Guangxuan Xiao, étudiant diplômé en génie électrique et informatique (EECS). et auteur principal d'un article sur StreamingLLM.
Les co-auteurs de Xiao incluent son conseiller, Song Han, professeur agrégé à l'EECS, membre du MIT-IBM Watson AI Lab et éminent scientifique de NVIDIA ; ainsi que Yuandong Tian, chercheur scientifique chez Meta AI ; Beidi Chen, professeur adjoint à l'Université Carnegie Mellon ; et l'auteur principal Mike Lewis, chercheur scientifique chez Meta AI. Les travaux seront présentés à la Conférence internationale sur les représentations d’apprentissage.
Un phénomène déroutant
Les grands modèles de langage codent les données, comme les mots dans une requête utilisateur, dans des représentations appelées jetons. De nombreux modèles utilisent ce que l'on appelle un mécanisme d'attention qui utilise ces jetons pour générer un nouveau texte.
En règle générale, un chatbot IA écrit un nouveau texte basé sur le texte qu'il vient de voir. Il stocke donc les jetons récents en mémoire, appelé cache KV, pour les utiliser ultérieurement. Le mécanisme d'attention construit une grille qui inclut tous les jetons dans le cache, une « carte d'attention » qui indique la relation entre chaque jeton, ou mot, les uns par rapport aux autres jetons.
Comprendre ces relations est une fonctionnalité qui permet à de grands modèles de langage de générer un texte de type humain.
Mais lorsque le cache devient très volumineux, la carte d’attention peut devenir encore plus volumineuse, ce qui ralentit le calcul.
De plus, si le codage du contenu nécessite plus de jetons que ce que le cache peut contenir, les performances du modèle diminuent. Par exemple, un modèle populaire peut stocker 4 096 jetons, alors qu’il existe environ 10 000 jetons dans un article universitaire.
Pour contourner ces problèmes, les chercheurs utilisent un « cache coulissant » qui supprime les jetons les plus anciens pour en ajouter de nouveaux. Cependant, les performances du modèle chutent souvent dès que le premier jeton est expulsé, réduisant rapidement la qualité des mots nouvellement générés.
Dans ce nouvel article, les chercheurs ont réalisé que s'ils conservaient le premier jeton dans le cache coulissant, le modèle conserverait ses performances même lorsque la taille du cache était dépassée.
Mais cela n'avait aucun sens. Le premier mot d’un roman n’a probablement rien à voir avec le dernier mot, alors pourquoi le premier mot serait-il si important pour que le modèle génère le mot le plus récent ?
Dans leur nouvel article, les chercheurs ont également découvert la cause de ce phénomène.
L’attention diminue
Certains modèles utilisent une opération Softmax dans leur mécanisme d'attention, qui attribue un score à chaque jeton qui représente sa relation avec les autres jetons. L'opération Softmax nécessite que tous les scores d'attention totalisent 1. Comme la plupart des jetons ne sont pas fortement liés, leurs scores d'attention sont très faibles. Le modèle supprime tout score d'attention restant dans le premier jeton.
Les chercheurs appellent ce premier jeton un « puits d’attention ».
«Nous avons besoin d'un récepteur d'attention, et le modèle décide d'utiliser le premier jeton comme récepteur d'attention car il est globalement visible – tous les autres jetons peuvent le voir. Nous avons constaté que nous devons toujours garder l'attention dans le cache pour maintenir la dynamique du modèle », explique Han.
En créant StreamingLLM, les chercheurs ont découvert que le fait d'avoir quatre jetons de récepteur d'attention au début du cache coulissant conduit à des performances optimales.
Ils ont également constaté que le codage positionnel de chaque jeton doit rester le même, même lorsque de nouveaux jetons sont ajoutés et que d’autres sont supprimés. Si le jeton 5 est supprimé, le jeton 6 doit rester codé comme 6, même s'il s'agit désormais du cinquième jeton dans le cache.
En combinant ces deux idées, ils ont permis à StreamingLLM de maintenir une conversation continue tout en surpassant une méthode populaire utilisant le recalcul.
Par exemple, lorsque le cache contient 256 jetons, la méthode de recalcul prend 63 millisecondes pour décoder un nouveau jeton, tandis que StreamingLLM prend 31 millisecondes. Cependant, si la taille du cache atteint 4 096 jetons, le recalcul nécessite 1 411 millisecondes pour un nouveau jeton, tandis que StreamingLLM n'a besoin que de 65 millisecondes.
« L'approche innovante de StreamingLLM, centrée sur le mécanisme de puits d'attention, garantit une utilisation et des performances stables de la mémoire, même lors du traitement de textes d'une longueur maximale de 4 millions de jetons », déclare Yang You, jeune professeur présidentiel d'informatique à l'Université nationale de Singapour, qui n'a pas participé à ces travaux. « Cette capacité n'est pas seulement impressionnante ; c'est transformateur, permettant à StreamingLLM d'être appliqué à un large éventail d'applications d'IA. Les performances et la polyvalence de StreamingLLM en font une technologie très prometteuse, sur le point de révolutionner la façon dont nous abordons les applications de génération basées sur l'IA.
Tianqi Chen, professeur adjoint aux départements d'apprentissage automatique et d'informatique de l'Université Carnegie Mellon qui n'a pas non plus participé à cette recherche, est d'accord, affirmant que « le streaming LLM permet d'étendre en douceur la durée de conversation de grands modèles de langage. Nous l’utilisons pour permettre le déploiement des modèles Mistral sur iPhone avec beaucoup de succès.
Les chercheurs ont également exploré l’utilisation de récepteurs d’attention lors de la formation du modèle en ajoutant plusieurs jetons d’espace réservé dans tous les échantillons de formation.
Ils ont constaté que l'entraînement avec des récepteurs d'attention permettait à un modèle de maintenir ses performances avec un seul récepteur d'attention dans son cache, plutôt que les quatre habituellement nécessaires pour stabiliser les performances d'un modèle pré-entraîné.
Mais même si StreamingLLM permet à un modèle de mener une conversation continue, le modèle ne peut pas se souvenir des mots qui ne sont pas stockés dans le cache. À l'avenir, les chercheurs prévoient de cibler cette limitation en étudiant des méthodes permettant de récupérer les jetons qui ont été expulsés ou de permettre au modèle de mémoriser des conversations précédentes.
StreamingLLM a été intégré à la grande bibliothèque d'optimisation de modèles de langage de NVIDIA, TensorRT-LLM.
Ce travail est financé en partie par le MIT-IBM Watson AI Lab, le MIT Science Hub et la National Science Foundation des États-Unis.
