Réorganisation des réseaux de mémoire à long terme : XLSTM pour l'IA de nouvelle génération

Les LSTM, ou réseaux de mémoire à long terme et à court terme, existent depuis longtemps. Ils ont été appliqués à de nombreuses tâches liées aux séquences, telles que la génération et la traduction de texte ou même la génération de légendes d'images.

Leur inconvénient était qu'ils ne pouvaient pas être parallélisés pour utiliser les puissants GPU modernes. Cette limitation a ouvert la voie à l’émergence de transformateurs exploitant les GPU pour une parallélisation massive de la formation et de l’inférence.

Si nous essayons maintenant de réorganiser et de paralléliser les LSTM, peuvent-ils devenir des outils pour construire des LLM de nouvelle génération ?

C'est exactement la question à laquelle répond l'article « XLSM — Extended Long Short-term Memory Networks », qui signifie « Extended » Long Short-Term Memory. Pour ce faire, ils proposent deux nouveaux blocs dans l'architecture, à savoir sLSTM et mLSTM.

Examinons donc en profondeur les blocs sLSTM et mLSTM proposés dans cet article et voyons comment nous pouvons les empiler ensemble pour développer l'architecture XLSTM.

Explication visuelle

Si vous êtes quelqu'un comme moi et que vous souhaitez expliquer visuellement les XLSTM, veuillez consulter la vidéo YouTube sur cet article :

Mise à jour LSTM

L'un des premiers réseaux conçus pour traiter les données séquentielles est le réseau neuronal récurrent.

Réseau neuronal récurrent

Il utilise des connexions récurrentes dans son architecture avec x comme entrée et o comme sortie. Si nous le déplions, nous pouvons le visualiser comme une séquence d'opérations se déroulant aux horodatages t-1, t et t+1 . Un inconvénient majeur des RNN était le problème du gradient de disparition, où le gradient atteint zéro lorsque nous empilons trop de blocs.

Les LSTM, ou réseaux de mémoire à long terme, ont été proposés pour surmonter les gradients de disparition en introduisant des états cellulaires et des mécanismes de déclenchement dans le réseau.

Une illustration simplifiée des LSTM

Les états cellulaires, c , sont des mémoires à long terme qui vivent sur plusieurs horodatages. Les états cachés, h , sont des mémoires à court terme qui se transmettent d'un pas de temps à un autre. Et bien sûr, nous avons les entrées, z , de la séquence d’entrée.

Trois portes ont des fonctions en forme de S. La porte d'oubli utilise une fonction sigmoïde pour décider quelles informations oublier dans la mémoire à long terme. La porte d'entrée utilise également une fonction sigmoïde pour traiter l'entrée et l'ajoute à la sortie de la porte d'oubli. Cette opération d'addition a un terme sophistiqué appelé carrousel d'erreurs constantes dans l'article XLSTM et dans la littérature académique. Cette opération d’addition permet de résoudre le problème des gradients de disparition rencontré dans les RNN. La sortie c_t est ensuite traitée par la porte de sortie , qui est généralement une fonction tanh menant à la sortie d'état caché h_t qui est transmise à l'étape suivante.

Avec ces opérations, nous avons décortiqué les deux équations principales des LSTM qui sont celle de c_t et h_t .

Inconvénient 1 – Réviser les décisions de stockage

L’un des principaux inconvénients des LSTM est leur incapacité à réviser les décisions de stockage. Cela signifie qu’à mesure que la longueur de la séquence augmente, le modèle devrait être capable de décider s’il conserve ou non les informations passées en mémoire.

Par exemple, si l’on prend cette phrase : « Tom est allé au magasin. Il a acheté des boissons », et comparez-le avec « Tom est allé au magasin pour faire l'épicerie, notamment des carottes, des oignons, des bananes, des pommes, des oranges, du café et du pain. Il a aussi acheté des boissons. Pour chaque nouveau mot, comme bananes ou pommes, le modèle doit constamment réviser s'il doit conserver le mot passé « Tom » dans sa mémoire. Il s’agit d’un grand défi pour les LSTM, et cela découle de la fonction sigmoïde de sa porte d’oubli.

La fonction sigmoïde vs la fonction exponentielle. Le sigmoïde s'aplatit vers les extrémités, mais l'exponentielle ne fait qu'augmenter.

Ainsi, si nous prenons une porte oubliée, elle est composée d’une fonction sigmoïde, qui présente une courbe en forme de S qui s’aplatit vers l’extrémité. Cela indique qu'à mesure que nous nous dirigeons vers des valeurs d'entrée plus élevées, la décision quant à ce qu'il faut oublier et ce qu'il faut garder en mémoire devient assez difficile. Mais si nous utilisons une fonction exponentielle à sa place, alors le jeu change, et à mesure que nous passons à des valeurs d'entrée plus élevées, nous obtenons une plage plus large pour les sorties. Ceci, à son tour, indique que les LSTM peuvent mieux réviser les décisions de stockage.

Solution 1 — sLSTM

Ainsi, la solution proposée dans cet article concerne les blocs système. Si l'on revient à l'équation LSTM classique qui représente l'état de la cellule, comme nous l'avons vu précédemment, elle est fonction de la porte d'oubli et des portes d'entrée.

Ces portes, à leur tour, sont composées de fonctions sigmoïdes. Alors, et si on remplaçait ces fonctions sigmoïdes par des fonctions exponentielles ? Les nouvelles portes f_t et i_t deviennent désormais exp(f_t) et exp(i_t), et c'est à peu près la principale modification pour créer le bloc sLSTM.

Contrairement à la fonction sigmoïde, qui contraint les entrées dans une plage fixe, les fonctions exponentielles ont tendance à prendre de la valeur à mesure que l'entrée augmente, et elles ne normalisent pas naturellement la sortie pour qu'elle se situe entre, disons, 0 et 1 comme la fonction sigmoïde. fonction.

Nous devons donc introduire un nouvel état de normalisation , qui est fonction des portes d’oubli et d’entrée. Nous pouvons le considérer comme une moyenne mobile d’une valeur de normalisation.

Nous utilisons les valeurs de normalisation calculées pour normaliser la sortie ou le nouvel état caché.

Alors que la normalisation prend en charge les états cachés, pour empêcher l'exponentielle de faire exploser les portes d'oubli et d'entrée, nous devons introduire un stabilisateur . Il se présente sous la forme de fonctions log pour contrer l’effet des exponentielles et introduire de la stabilité. Ainsi, l'état du stabilisateur est le maximum du journal de la sortie de la porte d'oubli et de la porte d'entrée. Nous soustrayons ces valeurs de stabilisation de l'entrée et oublions les portes pour les stabiliser.

Inconvénient 2 - Mémoire et parallélisation

Le deuxième inconvénient des LSTM est le manque de parallélisation. Les LSTM ont été conçus pour gérer des données séquentielles, ce qui signifie qu'ils ont besoin du résultat du traitement de l'entrée précédente dans la séquence pour traiter l'entrée actuelle dans la séquence. Cet inconvénient particulier empêche la parallélisation et est à l’origine de l’aube de l’ère Transformers.

La solution proposée dans cet article concerne les nouveaux blocs mLSTM. Alors, regardons-les ensuite.

Solution — mlSTM

Le prochain élément constitutif des XLSTM est le bloc mLSTM, où m représente la mémoire. Revenons à l'équation LSTM classique pour voir quel en est l'inconvénient. On voit que l’état de la cellule c_t est un scalaire. Cela signifie que nous ne traitons qu'un seul numéro à la fois lorsque nous avons le luxe de disposer de GPU modernes avec au moins 12 Go de mémoire.

Le bloc mLSTM introduit des matrices à la place des scalaires pour les états des cellules. Pour en revenir à notre équation classique de LSTM, et si nous remplacions le c_t par une matrice C *_t* afin que l'état de la cellule devienne maintenant C *_t* majuscule pour indiquer les matrices et que les états de la cellule puissent être récupérés non seulement par une porte i_t mais en stockant des paires clé-valeur qui sont des vecteurs. Dont les valeurs peuvent être récupérées par des requêtes qui sont des vecteurs de même dimension.

Pour que cela semble familier à la terminologie du transformateur, ils ont introduit ici la clé et la valeur pour former cette matrice.

XLSTM

Avec ces informations sur le sLSTM et le mLSTM, plongeons dans l'architecture détaillée de XLSTM.

sLSTM

Un aperçu détaillé du bloc sLSTM

En ce qui concerne sLSTM, nous utilisons des projections post-up. Ainsi, l'entrée passe d'abord à travers des couches de convolution causale avec une fonction d'activation swish. La sortie de ces couches est ensuite transmise à travers une couche linéaire de blocs diagonales avec quatre blocs diagonaux ou « têtes ». La sortie de ceux-ci est ensuite transmise via le bloc sLSTM à quatre têtes. Enfin, la sortie est projetée vers le haut à l'aide d'une couche MLP fermée avec activation GeLU et projetée vers le bas à l'aide d'une fonction MLP fermée.

mLSTM

Passant aux détails du bloc mLSTM, nous utilisons des projections préalables. Cela signifie que l'entrée est d'abord projetée vers le haut avec un facteur de projection de 2. L'une des sorties de projection va au mLSTM et l'autre à la porte de sortie. L'entrée dans le bloc mLSTM passe par une convolution causale, puis par des matrices de projection diagonale de bloc de taille de bloc 4, qui génèrent la requête, la clé et la valeur facilement utilisées par le bloc mLSTM.

Architecture XLSTM

Enfin, on peut ensuite empiler les deux types de blocs pour former l'architecture LSTM étendue. Ainsi, les blocs gris foncé sont les blocs mLSTM et les blocs gris clair sont les blocs sLSTM.

En termes d'avantages, l'article mentionne que les réseaux XLSTM ont une complexité de calcul linéaire et une complexité de mémoire constante concernant la longueur de séquence.

Évaluations

Les auteurs se sont formés sur l'ensemble de données SlimPajama pour le comparer à d'autres méthodes basées sur des transformateurs comme LAMA et des méthodes basées sur l'espace d'état comme MAMBA. Ils ont utilisé cette notation xLSTM a:b, où a est le nombre de blocs mLSTM et b est le nombre de blocs sLSTM dans la pile.

En termes de précisions, ils rapportent des précisions relatives en échelonnant les précisions entre 0 et 1, où 0 est aléatoire et 1 est parfait.

Les évaluations indiquent que XLSTM fonctionne mieux dans des tâches telles que Parity, tandis que Llama et Mamba fonctionnent mal.

D'après les résultats, la tâche de parité est particulièrement intéressante, dans laquelle les transformateurs ou les modèles d'espace d'état ont tendance à avoir des difficultés sans mélange de mémoire ni suivi d'état. Nous pouvons voir que dans ce type de tâche, le xLSTM atteint une précision de 1 lorsque nous utilisons à la fois les blocs sLSTM et mLSTM.

Ils ont également réalisé des études d’ablation pour montrer la robustesse des XLSTM. Ils sont faciles à comprendre à partir du document. De plus, cet article porte davantage sur les nouveautés architecturales des XLSTM, je n'entre donc pas ici dans les résultats expérimentaux.

Crier

Si vous avez aimé cet article, pourquoi ne pas me suivre Twitter où je partage les mises à jour des recherches des meilleurs laboratoires d'IA chaque jour de la semaine ?

Abonnez-vous également à mon Chaîne Youtube où j'explique visuellement les concepts et les articles sur l'IA.

Conclusion

J'espère que cet article a simplifié et facilité la compréhension de l'architecture XLSTM, pourquoi nous en avons besoin et comment elles pourraient potentiellement dépasser les transformateurs dans un avenir proche.

Attendons de voir ce qu'ils nous réservent. Je vous verrai dans mon prochain…