StreamingLLM - implémenter un modèle de langage en streaming efficace avec un attention sink
(github.com/mit-han-lab)- StreamingLLM est un framework permettant de déployer des LLM dans des applications de streaming à entrée de longueur infinie, sans sacrifier l’efficacité ni les performances
- Il traite les problèmes suivants : dans les dialogues multi-tours nécessitant de longues interactions, la mise en cache des états Key/Value des tokens précédents consomme beaucoup de mémoire, et les LLM classiques généralisent mal à des textes plus longs que la longueur de séquence vue à l’entraînement
- Les approches récentes de window attention qui ne mettent en cache que les KV échouent lorsque la longueur du texte dépasse la taille du cache ; StreamingLLM restaure largement les performances de la window attention grâce à un attention sink qui conserve les KV des tokens initiaux
- Il permet à des LLM entraînés avec une fenêtre d’attention de longueur finie de se généraliser à des séquences de longueur infinie sans fine-tuning, et réalise une modélisation du langage stable et efficace sur Llama-2, MPT, Falcon et Pythia jusqu’à plus de 4 millions de tokens
- En configuration de streaming, il montre jusqu’à 22,2× speedup par rapport à la baseline de sliding window recomputation
- Il n’étend pas la fenêtre de contexte, conserve seulement les tokens récents et l’attention sink, et supprime les tokens intermédiaires
- Si Llama-2 a été préentraîné avec une fenêtre de contexte de 4096 tokens, alors la taille maximale du cache de Llama-2 dans StreamingLLM est elle aussi de 4096
- Si on fournit un livre long en entrée, le modèle ne reconnaît que les tokens les plus récents et peut donc seulement résumer la partie de conclusion
- Les cas d’usage adaptés sont les applications de streaming qui doivent fonctionner en continu tout en évitant une dépendance aux données passées ou une forte consommation mémoire ; parmi les exemples cités figurent les dialogues multi-tours et les assistants quotidiens basés sur des LLM
- Cette approche est orthogonale aux méthodes récentes d’extension de contexte et peut être intégrée avec elles ; dans le cadre de StreamingLLM, l’extension de contexte signifie la possibilité de stocker davantage de tokens récents avec une taille de cache plus grande
- Exemple d’exécution :
examples/run_streaming_llama.py --enable_streaming; la configuration d’environnement utilise Python 3.8,torch,transformers==4.33.0,accelerate,datasets,evaluate,wandb,scikit-learn,scipy,sentencepiece - Le code principal a déjà été publié, notamment pour Llama-2, MPT, Falcon et Pythia ; le code d’évaluation de la perplexité et la démo Streaming Llama Chatbot sont également disponibles, tandis que le dataset StreamEval et son code d’évaluation ne sont pas encore publiés
1 commentaires
Commentaires Hacker News
On dirait que certains prennent ce travail pour quelque chose comme de l’attention dense complète.
Ici, il ne s’agit pas de permettre au modèle de voir du contenu très éloigné, mais plutôt d’une amélioration d’efficacité qui réutilise le cache pour maintenir la perplexité, au lieu de recalculer à chaque fois une fenêtre glissante et de payer un coût en L² répété T fois.
Les tests ont aussi été mesurés en concaténant quelque chose comme
Q A Q A Q A Q A..., et non en demandant de retrouver une réponse beaucoup plus loin dans un schéma du typeQ Q Q Q A A A A....Mesurer la perplexité signifie produire un « texte lisible », c’est-à-dire des phrases localement plausibles, pas prouver qu’on « extrait » quelque chose depuis un grand trou triangulaire hors de portée de l’attention.
Si on lui donnait un livre en demandant d’écrire le premier mot de chaque paragraphe, ou de résumer chaque chapitre en une phrase, je pense qu’il échouerait.
Après l’avoir testé moi-même, il ne m’a pas semblé que cela permette une extension de la longueur de contexte, et l’exécution était plutôt rapide.
Il utilisait environ 35 Go de mémoire sur A100, avec une consommation restée fixe pendant toute l’exécution.
J’ai pris un livre sur Project Gutenberg, je l’ai découpé en paragraphes, je les ai fournis un par un en lui demandant de répondre « okay » à chaque paragraphe, puis j’ai posé une question à la fin : la réponse était une hallucination complète.
À noter que, pendant les quelque 10 minutes où je l’ai manipulé, je n’ai même pas vraiment réussi à faire répondre en anglais le modèle par défaut,
lmsys/vicuna-13b-v1.3.https://gist.github.com/bluecoconut/9cae9e91fe3b1616ed650a96...
Cela reste un travail intéressant, et le point central semble être la découverte de la Figure 2.
Les deux premières couches montrent un motif local, avec davantage d’attention portée aux tokens récents, mais une fois passées ces couches, le modèle porte une forte attention aux tokens initiaux dans toutes les couches et toutes les têtes.
Les auteurs appellent cela des « attention sinks » ; même si ces tokens ne sont pas importants sémantiquement, la somme des scores d’attention doit valoir 1 à cause du Softmax, donc l’attention restante doit bien aller quelque part.
Leur explication est que, dans un modèle de langage autorégressif, les tokens initiaux sont visibles depuis presque tous les tokens suivants, ce qui les rend faciles à apprendre comme puits.
StreamingLLM ressemble davantage à un « hack » qui corrige ce comportement étrange lorsque l’on découpe la fenêtre d’attention d’un LLM, et c’est un cas qui révèle une fissure dans l’usage du Softmax ; si l’on veut des LLM flexibles sur la longueur de contexte, une autre fonction pourrait peut-être être préférable.
Au premier survol, cela paraît presque trop beau pour être vrai, mais la qualité du travail semble correcte et la technique est étonnamment simple.
L’idée consiste à appliquer l’attention, dans chaque couche, uniquement au premier token et à la fenêtre de contexte glissante, en ignorant les tokens intermédiaires.
Cela semble vouloir dire que chaque couche pousse peu à peu les informations pertinentes vers l’arrière de la séquence, afin que la fenêtre d’attention glissante à la fin des couches supérieures puisse les voir.
Mais si la portée couverte par toutes les fenêtres glissantes n’est pas suffisante pour relier toute la séquence, il se peut que toutes les informations importantes ne puissent pas être transmises vers l’avant.
Par exemple, si toutes les fenêtres ont la même longueur et que
profondeur du modèle × longueur de fenêtre < longueur de séquence, une limite apparaît.Cela semble rendu possible par l’observation que Softmax doit faire en sorte que la somme vaille 1.
À première vue, le modèle a tendance à utiliser le premier token comme placeholder lorsqu’il n’a pas besoin de prêter attention aux tokens précédents.
La première fois que j’ai vu ce problème, c’était dans un post HN d’Evan Miller, qui expliquait qu’il était erroné de forcer les têtes d’attention à distribuer toute leur attention entre les tokens précédents, et qu’il fallait ajouter 1 au dénominateur du Softmax pour permettre de « ne pas prêter attention ».
C’est bien d’exploiter cette observation sans réentraînement, et je me demande aussi comment le modèle aurait changé si l’on avait suivi la proposition d’Evan.
[2] https://news.ycombinator.com/item?id=36851494
Ils ont entraîné le modèle avec un token puits dédié dont toutes les valeurs valent 0, mais même ainsi, d’autres tokens initiaux finissent par être utilisés comme puits ; leur conclusion semble donc être qu’il vaut mieux avoir un token puits dédié.
À ma connaissance, cela n’a jamais vraiment amélioré les performances.
Pour manipuler la fenêtre d’attention après l’entraînement, Softmax + 1 pourrait être mieux adapté, mais je ne sais pas si quelqu’un l’a testé à grande échelle.
Ajouter une mémoire de cache d’attention est une solution très intéressante à ce problème.
Il y a quelques jours, un article est aussi sorti avec une observation apparentée sur les Vision Transformers.
Les modèles Transformer semblent choisir des tokens où stocker de l’information globale, et paraissent avoir besoin d’une sorte de « token de pensée ».
Fournir des tokens spécifiques à cet usage améliore un peu les performances, et les visualisations explicatives obtenues sont assez intéressantes.
[0] https://arxiv.org/pdf/2309.16588.pdf
En fine-tuning, on pourrait figer les paramètres du modèle original et n’ajuster que les paramètres entrant et sortant de ces nouvelles unités de cache de « tuning ».
On pourrait alors remplacer ou combiner différents ensembles d’unités de tuning.
Par exemple, une unité d’évitement des grossièretés + une unité de terminologie spécifique + une unité pour écrire de manière concise : une sorte de super-prompt composé de plusieurs éléments.
Si le nombre de nouveaux paramètres est suffisamment faible, cela consommerait plus de mémoire, mais un optimiseur d’ordre élevé pourrait permettre un tuning rapide et efficace.
On peut aussi imaginer augmenter simultanément la longueur des séquences et le nombre d’unités pendant l’entraînement.
Par exemple, n’utiliser que quelques unités pour les séquences courtes, puis ajouter des unités au fur et à mesure que la longueur des séquences d’entraînement augmente, en poursuivant l’entraînement.
Au lieu d’un calendrier arbitraire, on pourrait aussi contrôler l’extension du cache à partir des performances ou d’une analyse des gradients.
Les auteurs ont publié une FAQ, qui peut dissiper une partie de la confusion : https://github.com/mit-han-lab/streaming-llm/blob/main/READM...
À la question « Peut-on mettre un long texte, comme un livre, dans StreamingLLM pour le résumer ? », ils répondent qu’on peut bien fournir un long texte, mais que le modèle ne perçoit que les tokens les plus récents ; donc, si on lui donne un livre, il ne pourra résumer que les derniers paragraphes, ce qui risque de ne pas être très utile.
Autrement dit, il ne s’agit pas d’étendre la fenêtre de contexte du LLM ni de renforcer sa mémoire à long terme ; le point fort de StreamingLLM est de générer un texte fluide à partir des tokens récents, sans rafraîchir le cache.
Je peux me tromper, mais je n’ai pas l’impression que ce soit une technique permettant à un LLM de se référer à du contenu au-delà de la longueur sur laquelle il a été entraîné, comme beaucoup semblent le penser.
Cela semble plutôt concerner le maintien des performances du modèle sur de longs textes, plus précisément sur le contenu qui se trouve encore dans la fenêtre de contexte.
L’explication est que le modèle apprend à placer une sorte de charge dans l’attention portée aux premiers tokens du texte, et que tout casse quand ces tokens sortent de la fenêtre ; mais je ne sais pas vraiment pourquoi.
Si ce n’est pas une entrée d’instruction, j’aurais tendance à penser que le texte du milieu est aussi bon que le texte du début.
Je me demande comment ces techniques de fenêtre glissante gèrent les cas où une instruction inattendue n’apparaît qu’à la fin.
Par exemple, supposons qu’on donne un livre au modèle et que la dernière phrase soit : « renvoie le nombre de lettres m dans l’entrée précédente ». Un humain soupirerait, relirait et compterait, mais un LLM n’a pas la capacité de revenir en arrière pour relire l’entrée.
Dans cet exemple, même en ignorant les limites propres aux LLM pour compter des caractères, il me semble que pour vraiment résoudre le problème, il faudrait que le LLM puisse faire des boucles et sauter arbitrairement dans le texte.
Bien sûr, cela créerait alors des problèmes entièrement nouveaux, et pourrait même nécessiter une nouvelle architecture.
L’idée serait de transformer la question en mots-clés pertinents, de chercher, puis de reparcourir le contenu pour trouver les informations utiles.
S’il existe un prétraitement nécessaire, un LLM capable d’« aller faire suffisamment de recherches puis répondre » pourrait être extrêmement puissant.
Depuis environ 10 000 ans, nous améliorons les techniques de gestion des connaissances pour dépasser les capacités et le temps du cerveau individuel ; les modèles de langage devraient donc exploiter de vraies méthodes de recherche et de prédigestion, pas seulement une recherche Bing.
La mémoire à court terme n’a pas besoin de se souvenir de ce que faisait tel ou tel bout de code : il suffit de l’étiqueter à la lecture puis de s’appuyer sur un index partagé de tags, extensible.
Cela dit, plus j’y pense, plus cela ressemble au préentraînement général des LLM, et plus l’index de connaissances me donne l’impression d’être un énorme amas de poids de LLM.
Ce serait davantage une couche posée au-dessus du LLM qu’un changement dans le fonctionnement du LLM lui-même.
Je comprends peut-être mal, mais dans ce cas, est-ce que l’état caché ne résout pas le problème de récupération ?
Comme il faut de toute façon absorber toute l’entrée avant de répondre, le fait que l’instruction soit au début ou à la fin ne semble pas avoir beaucoup d’effet en dehors de l’attention.
Claude 100K demande aux utilisateurs de placer l’instruction à la fin.
Sinon, on peut utiliser un modèle rapide pour vérifier s’il y a une instruction à la fin et la déplacer au début.
Je ne sais pas vraiment si les modèles actuels peuvent déjà faire cela avec des entrées courtes.
Pour le dire avec une pointe d’humour, les LLM travaillent vraiment très dur à réinventer les RNN, et si on leur donne les outils, ils finiront sans doute par y arriver.
Vu autrement, les modèles Transformer essaient de prédire quelles parties d’un réseau RNN « valent la peine d’être conservées » sous contrainte de ressources.
Les Transformers actuels utilisent une heuristique simple, et ce résultat améliore cette heuristique.
Comme pour beaucoup de problèmes NP-complets, même si l’on ne peut pas obtenir une solution parfaite, il peut exister des approximations utiles ; les Transformers montrent que c’est aussi possible avec les réseaux de neurones.
Il s’est situé pendant un temps vers le milieu des classements open source, donc c’est une approche assez légitime ; elle n’est simplement pas à la mode.
[1] : https://huggingface.co/blog/rwkv
Le principal avantage des Transformers sur les RNN est la parallélisation de l’entraînement.
Les RNN souffrent de disparition du gradient pendant l’entraînement, et il est aussi difficile d’obtenir un taux d’utilisation global élevé, car il faut de grands batchs ; ils sont donc délicats à gérer.
L’existence de modèles comme RWKV montre qu’il pourrait y avoir un avenir où l’on entraîne comme un Transformer et où l’on infère comme un RNN.
À ce propos, le professeur Han du MIT donne actuellement un cours TinyML ouvert au public.
https://news.ycombinator.com/item?id=37620507
https://efficientml.ai