Revue du transformateur avec une explication approfondie des principes de l'alternative Mamba (illustrée)

Une alternative au transformateur dans la modélisation linguistique

Transformateur L'architecture est un élément clé du succès des grands modèles linguistiques (LLM). Presque tous les grands modèles de langage utilisés aujourd'hui font appel à l'architecture, qu'il s'agisse de modèles à source ouverte tels que Mistral à des modèles à source fermée tels que ChatGPT.

Afin d'améliorer encore le grand modèle linguistique, de nouvelles architectures ont été développées, qui peuvent même aller au-delà de l'architecture Transformer. L'une de ces approches est le Mambaun type de modèle espace-état.

Mamba dans le journal Mamba : modélisation de séquences temporelles linéaires avec des espaces d'état sélectifs1 dans lequel il a été présenté. Pour en savoir plus, consultez son dépôt Retrouvez sa mise en œuvre officielle et les points de contrôle du modèle dans.

Dans ce billet, je présenterai le domaine de la modélisation de l'espace d'état dans le contexte de la modélisation linguistique et j'explorerai pas à pas divers concepts pour aider à comprendre le domaine. Nous discuterons ensuite de la manière dont Mamba pourrait remettre en question l'architecture Transformer.

En tant que guide visuel, cet article présentera de nombreuses visualisations pour aider à comprendre Mamba et les modèles d'espace d'état !

Première partie :Questions relatives aux transformateurs

Pour illustrer l'intérêt de l'architecture de Mamba, examinons d'abord brièvement Transformers et explorons l'une de ses lacunes.

Le transformateur traite toute entrée de texte comme s'il s'agissait d'une entrée de texte provenant de l'application jetons constituant séquence.

L'un des principaux avantages des transformateurs est que, quelle que soit l'entrée qu'il reçoit, il peut revenir à n'importe quel jeton antérieur de la séquence pour dériver sa représentation.

Composants principaux des transformateurs

Rappelez-vous que le transformateur se compose de deux structures, un ensemble de blocs d'encodage pour représenter le texte et un ensemble de blocs de décodage pour générer le texte. Combinées, ces structures peuvent être utilisées pour de multiples tâches, y compris la traduction.

Nous pouvons prendre cette structure et créer des modèles génératifs en utilisant uniquement des décodeurs. Ce modèle basé sur un transformateur, leTransformateur de pré-entraînement génératif(GPT) qui utilise des blocs décodeurs pour compléter une partie du texte d'entrée.

Voyons comment cela fonctionne !

Bénédiction dans la formation...

Un bloc décodeur unique se compose de deux éléments principaux : le mécanisme d'auto-attention masquée et le réseau neuronal de type feedforward.

Le mécanisme d'auto-attention est l'une des principales raisons pour lesquelles ces modèles fonctionnent si bien. Il permet d'obtenir une vue non comprimée de l'ensemble de la séquence et s'entraîne rapidement.

Comment cela fonctionne-t-il ?

Il crée une matrice qui prend chaque Jeton Comparer avec chaque token précédent. Les poids dans la matrice dépendent de la pertinence des paires de tokens l'une par rapport à l'autre.

Cette matrice est créée en une seule fois lors de la formation. "Mon"et"nom"L'attention entre n'a pas besoin d'être calculée avant d'être calculée.nom"et"est"Attention entre.

Elle réaliseparallélisationqui accélère considérablement la formation !

Problèmes de raisonnement !

Il y a cependant un inconvénient. Lors de la génération du jeton suivant, nous devons recalculer la valeur duséquence complètemême si nous avons déjà généré des jetons.

Générer la longueur de laLLa séquence nécessite environL²ce qui peut s'avérer coûteux en termes de calcul si la longueur de la séquence augmente.

Cette nécessité de recalculer l'ensemble de la séquence est l'un des principaux goulets d'étranglement de l'architecture Transformer.

Voyons comment une technique "classique", les réseaux neuronaux récurrents (RNN), résout ce problème d'inférence lente.

Le RNN est-il la solution ?

Le réseau neuronal récurrent (RNN) est un réseau basé sur la séquence. Il reçoit deux entrées à chaque pas de temps, c'est-à-dire le pas de tempstet le pas de temps précédentt-1de l'état caché est utilisée pour générer l'état caché suivant et prédire la sortie.

Le RNN dispose d'un mécanisme de boucle qui lui permet de transmettre des informations d'une étape à l'autre. Nous pouvons "étendre" ce processus de visualisation pour le rendre plus explicite.

Lors de la génération de la sortie, le RNN ne doit prendre en compte que l'état caché précédent et l'entrée actuelle. Cela évite le problème de recalculer tous les états cachés précédents requis par le transformateur.

En d'autres termes, un RNN est capable d'effectuer une inférence rapide parce qu'il s'adapte linéairement à la longueur de la séquence ! En théorie, il pourrait même avoirLongueur de contexte illimitée.

Pour illustrer cela, appliquons le RNN au texte d'entrée que nous avons utilisé précédemment.

Chaque état caché est une agrégation de tous les états cachés précédents, généralement une vue comprimée.

Il y a toutefois un problème à ce niveau : ......

Notez que lors de la génération du nom "Maarten"lorsque le dernier état caché ne contient plus d'informations sur le mot "Bonjour"Les RNN ont tendance à oublier des informations au fil du temps parce qu'ils ne prennent en compte que le dernier état.

Bien que les RNN soient plus rapides en termes de formation et d'inférence, ils n'ont pas la précision que les modèles Transformer peuvent fournir.

C'est pourquoi nous étudions les modèles d'espace d'état pour utiliser efficacement les RNN (et parfois la convolution).

Partie 2 :Modélisation de l'espace d'état (SSM)

Les modèles d'espace d'état (MSE), comme les transformateurs et les RNN, traitent des séquences d'informations, telles que le texte et les signaux. Dans cette section, nous présentons les concepts de base des modèles d'espace d'état et leur relation avec les données textuelles.

Qu'est-ce que l'espace d'état ?

Un espace d'états contient le nombre minimum de variables nécessaires pour décrire complètement un système. C'est une façon de représenter mathématiquement un problème en définissant les états possibles d'un système.

Simplifions. Imaginons que nous nous déplacions dans un labyrinthe. "espace d'état"Il s'agit d'une carte de tous les emplacements (états) possibles. Chaque point représente un emplacement unique dans le labyrinthe avec des détails spécifiques, comme la distance qui vous sépare de la sortie.

"représentation de l'espace d'état"est une description simplifiée de cette carte. Elle indique votre position actuelle (état actuel), la prochaine position vers laquelle vous pouvez vous rendre (état futur possible) et la façon dont vous pouvez vous rendre à la prochaine position (en vous déplaçant vers la droite ou la gauche).

Bien que le modèle de l'espace d'état utilise des équations et des matrices pour suivre ce comportement, il ne s'agit en fait que d'un moyen de savoir où l'on se trouve, où l'on peut aller et comment y parvenir.

Les variables décrivent un état, dans notre cas les coordonnées X et Y, et la distance à la sortie, qui peut être exprimée par "vecteur d'état".

Cela vous semble familier ? C'est parce que les embeddings ou les vecteurs sont également souvent utilisés dans les modèles de langage pour décrire l'"état" d'une séquence d'entrée. Par exemple, un vecteur décrivant votre position actuelle (un vecteur d'état) pourrait ressembler à ceci :

Dans les réseaux neuronaux, le terme "état" fait généralement référence à l'état caché du système, qui est l'un des aspects les plus importants de la génération de nouveaux tokens dans un grand modèle linguistique.

Qu'est-ce qu'un modèle d'espace d'état ?

Les modèles d'espace d'état sont une classe de modèles utilisés pour décrire ces représentations d'état et prédire leur état suivant, les prédictions étant basées sur certaines entrées.

Traditionnellement, dans le temps tSSM :

• La séquence d'entrée sera la suivante x(t)(par exemple, se déplacer de gauche à droite dans un labyrinthe) correspond à la représentation d'un état potentiel h(t)(par exemple, distance par rapport à la sortie et coordonnées x/y)
• et d'en déduire la séquence de sortie prédite y(t)(par exemple, se déplacer à nouveau vers la gauche pour atteindre la sortie plus rapidement)

Toutefois, au lieu d'utiliser des séquences discrètes (par exemple, en se décalant une fois vers la gauche), le SSM accepte des séquences continues en entrée et prédit la séquence de sortie.

Le SSM part du principe que les systèmes dynamiques (par exemple, les objets se déplaçant dans l'espace 3D) peuvent être temporellement t de l'état et deux équations pour le prédire.

En résolvant ces équations, nous supposons que nous pouvons révéler des principes statistiques permettant de prédire l'état du système sur la base des données observées (séquences d'entrée et états précédents).

L'objectif est de trouver cette représentation de l'état h(t)qui nous permet de passer des séquences d'entrée aux séquences de sortie.

Ces deux équations sont au cœur du modèle de l'espace d'état.

Ces deux équations seront référencées tout au long du guide. Pour les rendre plus intuitives, l'expressionNous utilisons un code couleurafin de pouvoir les citer rapidement.

équation d'état décrit comment l'état sera modifié en fonction des entrées (par l'intermédiaire de la fonction Matrice B) l'état de l'impact (via Matrice A) et le changement.

Comme indiqué précédemment.h(t) se réfère à un moment donné t La représentation de l'état potentiel dux(t) se réfère à une certaine entrée.

équation de sortie décrit comment l'état est transmis par le Matrice C converti en sortie, et comment les entrées sont converties en sorties par les Matrice D Affecte la production.

prendre note: Matrix A,B,C répondre en chantant D On parle aussi souvent de paramètresparce qu'elles peuvent être apprises.

En visualisant ces deux équations, on obtient l'architecture suivante :

Voyons étape par étape comment ces matrices affectent le processus d'apprentissage.

Supposons que nous disposions de signaux d'entrée x(t)Le signal est d'abord connecté au Matrice B Multiplier.Matrice B Décrit comment les intrants affectent le système.

L'état actualisé (semblable à l'état caché d'un réseau neuronal) est un espace potentiel qui contient la "connaissance" de base de l'environnement. Nous établissons un lien entre l'état et le Matrice A Multipliée, cette matrice décrit les associations entre tous les états internes et représente la dynamique sous-jacente du système.

Comme vous l'avez peut-être remarqué.Matrice A Appliquer avant la création de la représentation de l'état et à nouveau après la mise à jour de la représentation de l'état.

Ensuite, nous utilisons le Matrice C pour décrire comment les états sont convertis en sorties.

Enfin, nous pouvons utiliser la fonction Matrice D Fournit un signal direct de l'entrée à la sortie. On parle aussi couramment de connexion de saut.

en raison de Matrice D Tout comme les connexions par saut, les MSS sont souvent considérés comme les formes suivantes qui ne contiennent pas de connexions par saut.

En revenant à notre perspective simplifiée, nous pouvons maintenant nous concentrer sur la matrice A,B répondre en chantant C comme le cœur du SSM.

Nous pouvons mettre à jour les équations originales (et ajouter de jolies couleurs) pour indiquer à quoi sert chaque matrice, comme nous l'avons fait précédemment.

Ensemble, ces deux séries d'équations visent à prédire l'état du système à partir des données observées. Étant donné que les entrées sont continues, la représentation principale de la MSS est la suivante échelle de temps continue.

Des signaux continus aux signaux discrets

Si vous avez un signal continu, trouvez la représentation de l'état h(t) difficile à analyser. En outre, étant donné que nous disposons généralement d'entrées discrètes (par exemple, des séquences de texte), nous aimerions discrétiser le modèle.

Pour ce faire, nous utilisons la fonction technique de maintien de l'ordre zéro. Il fonctionne comme suit : tout d'abord, chaque fois que nous recevons un signal discret, nous conservons sa valeur jusqu'à ce que nous recevions un nouveau signal discret. Ce processus génère un signal continu pour le SSM :

Nous conservons la valeur du temps par un nouveau paramètre pouvant être appris stimulateur cardiaque ∆ Indique. Il indique la résolution de l'entrée.

Maintenant que nous avons généré un signal continu pour l'entrée, nous pouvons générer une sortie continue et échantillonner ces valeurs en fonction du pas de temps de l'entrée.

Ces valeurs échantillonnées sont le résultat de notre discrétisation !

Mathématiquement, nous pouvons appliquer la règle de l'ordre zéro de la manière suivante :

Ensemble, ils nous permettent de passer d'un MSS continu à un MSS discret. à ce stade, le modèle n'est plus un MSS discret. de fonction à fonction x(t) → y(t)Au lieu de cela, c'est de séquence à série xₖ → y_ₖ: :

Ici, la matrice A répondre en chantant B Désignons maintenant les paramètres discrétisés du modèle.

Nous utilisons k au lieu de tpour distinguer les SSM continus des SSM discrets.

Attention : Pendant la formation, nous conservons Matrice A dans sa forme continue, plutôt que dans sa version discrétisée. Au cours de la formation, la représentation continue est discrétisée.

Maintenant que nous disposons d'une formule pour la représentation discrétisée, examinons la manière de procéder pour obtenir une représentation discrétisée. compter Le modèle.

représentation récursive

Notre SSM discrétisé nous permet de construire le problème dans un pas de temps spécifique, plutôt que dans un signal continu. Comme nous l'avons vu précédemment pour les RNN, les méthodes récursives sont très utiles ici.

Si nous considérons des pas de temps discrets au lieu de signaux continus, nous pouvons reformuler le problème en utilisant des pas de temps :

À chaque pas de temps, nous calculons l'entrée actuelle (Bxₖ) comment il affecte l'état précédent (Ahₖ₋₁) et calcule ensuite la sortie prédite (Chₖ).

Cette représentation vous est peut-être déjà familière ! Nous pouvons l'analyser comme nous l'avons fait avec le traitement précédent des RNN.

Nous pouvons développer (ou développer en une série de pas de temps) comme suit :

Il convient de noter que nous pouvons utiliser une version discrétisée de cette méthode en utilisant l'approche de base d'un RNN.

représentation par convolution (math.)

Nous pouvons utiliser la convolution pour représenter le SSM. Rappelez-vous que dans les tâches classiques de reconnaissance d'images, nous appliquons des filtres (noyaux) pour extraire les caractéristiques agrégées :

Comme il s'agit de texte et non d'images, nous devons utiliser une vue unidimensionnelle :

Le noyau que nous utilisons pour représenter ce "filtre" est dérivé de la formule SSM :

Voyons ce que ce noyau fait en pratique. Comme pour la convolution, nous pouvons utiliser le noyau SSM pour parcourir chaque ensemble de jetons et calculer la sortie :

Cela montre également l'effet que le remplissage peut avoir sur le résultat. J'ai modifié l'ordre du remplissage pour améliorer la visualisation, mais nous appliquons généralement le remplissage à la fin de la phrase.

Lors de l'étape suivante, le noyau se déplace une fois pour effectuer l'étape suivante du calcul :

Dans la dernière étape, nous pouvons voir l'effet complet du noyau :

L'un des principaux avantages de la représentation des MSS sous forme de convolutions est qu'ils peuvent être formés en parallèle comme les réseaux neuronaux convolutifs (CNN). Cependant, en raison de la taille fixe du noyau, leur inférence n'est pas aussi rapide et sans contrainte que les RNN.

Trois types de représentation

Ces trois représentations -progression,récursif (calcul) répondre en chantant convolution Chacun d'entre eux présente des avantages et des inconvénients différents :

Il est intéressant de noter que nous pouvons désormais utiliser des MSS récursifs pour une inférence efficace tout en utilisant des MSS convolutifs pour un apprentissage parallèle.

Grâce à ces représentations, nous pouvons utiliser une astuce astucieuse consistant à choisir une représentation en fonction de la tâche. Pendant l'apprentissage, nous utilisons une représentation convolutive qui peut être parallélisée, tandis que pendant l'inférence, nous utilisons une représentation récursive efficace :

Ce modèle est connu sous le nom de Couche d'espace d'état linéaire (LSSL). 2

Ces représentations partagent une propriété importante, à savoirinvariant linéaire en temps (math.)（LTI indique les paramètres du SSM. A,B répondre en chantant C est fixe à tous les pas de temps. Cela signifie que la matrice A,B répondre en chantant C est le même pour chaque jeton généré.

En d'autres termes, quelle que soit la séquence que vous donnez au SSM.A,B répondre en chantant C Les valeurs de chacun d'entre eux restent les mêmes. Nous avons une représentation statique qui ne se soucie pas du contenu.

Avant de voir comment Mamba résout ce problème, examinons la dernière pièce du puzzle - le système de gestion de l'information.Matrice A.

matrices A l'importance de

L'un des aspects les plus importants de la formule du SSM est sans aucun doute le fait que l'Union européenne est le seul État membre à disposer d'un mécanisme de contrôle de la qualité. Matrice A. Comme nous l'avons vu précédemment dans la représentation récursive, elle capture des informations sur la précédent des informations sur l'État pour construire le renouvelé Statut.

Essentiellement.Matrice A Générer un état caché :

Par conséquent, la création de Matrice A peut déterminer si nous pouvons nous souvenir de quelques tokens précédents ou si nous pouvons saisir tous les tokens que nous avons vus jusqu'à présent. en particulier dans le contexte de la représentation récursive, qui n'est possible qu'avec des rappelle que État précédent.

Comment créer de manière à conserver une grande partie de la mémoire (taille du contexte) ? matrice A?

Nous utilisons Hungry Hungry Hippo ! ou HiPPO3 réaliservotre (honorifique)escaliers comment (dans quelle mesure)terme (math.) se mettre entre les mains defilm chanceoperator.HiPPO tente de compresser tous les signaux d'entrée qu'il a vus jusqu'à présent en un vecteur de coefficients.

Il utilise matrice A pour construire une représentation d'état qui capture mieux les jetons récents et diminue les jetons plus anciens. La formule peut être exprimée comme suit :

Supposons que nous ayons un carré matrice ACela nous donne :

Construire avec HiPPO matrice A s'est avérée bien meilleure que l'initialisation à une matrice aléatoire. Il est donc plus performant pour reconstruire la mise à jour L'aspect de la signalisation (jeton le plus proche) est plus important que l'aspect de l'information. plus ancien Le signal (jeton initial) est plus précis.

L'idée centrale de la matrice HiPPO est de générer un état caché qui se souvient de son histoire.

Mathématiquement, il le fait en suivant Polynôme de Legendre pour y parvenir, ce qui lui permet d'approcher tous les enregistrements historiques.4

La méthode HiPPO est ensuite appliquée aux représentations récursives et convolutives que nous avons vues précédemment pour traiter les dépendances à longue distance. Le résultat. Espace d'état structuré pour les séquences (S4)La classe SSM est une classe de SSM qui peut traiter efficacement de longues séquences.5

Il se compose de trois parties :

• modèle espace-état
• HiPPO pour la transformation dépendance à longue distance
• La discrétion est utilisée pour créer récursif (calcul) répondre en chantant convolution indiquer

Ce type de SSM présente plusieurs avantages, en fonction de la représentation choisie (récursive ou convolutive). Il peut également traiter de longues séquences de texte et stocker des mémoires de manière efficace en étant basé sur des matrices HiPPO.

prendre noteSi vous souhaitez mieux comprendre comment calculer la matrice HiPPO et construire votre propre modèle S4, je vous recommande vivement de lire le document suivant S4 annoté.

Troisième partie :Mamba - un modèle sélectif de l'espace d'état

Nous avons finalement couvert toutes les bases nécessaires pour comprendre ce qui fait la spécificité de Mamba. Les modèles d'espace d'état peuvent être utilisés pour modéliser des séquences de texte, mais il y a encore un certain nombre d'inconvénients que nous aimerions éviter.

Dans cette section, nous examinerons les deux principales contributions de Mamba :

1. un type de Algorithme de balayage sélectifpermettant au modèle de filtrer les informations (non) pertinentes
2. un type de Algorithmes tenant compte du matérielAu moyen de l'outil balayage parallèle,fusion des noyaux répondre en chantant recalculer pour stocker efficacement les résultats (intermédiaires).

Ensemble, ils créent le Modélisation sélective de l'espace d'état peut-être S6 modèle qui peut être utilisé comme l'auto-attention pour créer des Bloc Mamba.

Avant d'examiner ces deux principales contributions, voyons d'abord pourquoi elles sont nécessaires.

Quel est le problème que l'on tente de résoudre ?

Les modèles d'espace d'état, même les modèles S4 (modèles d'espace d'état structurés), sont peu performants pour certaines tâches clés de la modélisation et de la génération de langues, à savoir Capacité à se concentrer sur des données spécifiques ou à les ignorer.

Nous pouvons illustrer cela avec deux tâches synthétiques, à savoir copie sélective répondre en chantant tête de l'inducteur.

existent copie sélective L'objectif du SSM dans une tâche est de copier les éléments d'entrée et de les restituer dans l'ordre :

Toutefois, étant donné que le SSM est Invariant en temps linéaire(utilisé comme expression nominale)qu'il est peu performant dans cette tâche. Comme nous l'avons vu précédemment, la matrice A,B répondre en chantant C Chaque jeton généré est le même pour les SSM.

Par conséquent, le SSM n'est pas en mesure d'exécuter les tâches suivantes Raisonnement en fonction du contenuC'est un problème parce que nous voulons que SSM raisonne sur les entrées (hints). parce qu'il traite chaque jeton de la même manière en raison d'une matrice fixe de A, B et C. C'est un problème parce que nous voulons que SSM raisonne sur les entrées (hints).

Les performances du SSM sont médiocres dans une autre tâche, celle de l'évaluation de la qualité de l'eau. tête de l'inducteurdont l'objectif est de reproduire les schémas trouvés dans les données d'entrée :

Dans l'exemple ci-dessus, il s'agit essentiellement d'une invite unique dans laquelle nous essayons d'"enseigner" au modèle à utiliser le même indice dans chaque ".Q." après avoir fourni un "A.". Cependant, comme le SSM est invariant dans le temps, il ne peut pas choisir les jetons précédents à rappeler dans l'historique.

Pour ce faire, nous nous concentrerons sur Matrice B pour illustrer ce point. Indépendamment de l'entrée x Ce que c'est.Matrice B est toujours le même, donc avec x Sans objet :

De même.A répondre en chantant C reste également fixe à tout moment, indépendamment de l'entrée. Cela suggère que ce que nous avons vu jusqu'à présent du système SSM état de non-fonctionnement Caractéristiques.

En revanche, ces tâches sont relativement simples pour le transformateur, puisqu'elles sont basées sur la séquence d'entrée dynamique (science) Changement d'attention. Ils peuvent sélectivement "regarder" ou "se concentrer" sur différentes parties de la séquence.

La faible performance de la GSS dans ces tâches illustre les problèmes potentiels de la GSS invariante dans le temps, de la matrice A,B répondre en chantant C La nature statique de la Sens du contenu du problème.

Rétention sélective des informations

La représentation récursive du SSM crée un état plus petit qui est très efficace parce qu'il comprime l'ensemble de l'historique. Cependant, comparé au modèle Transformer, le modèle Transformer ne compresse pas l'historique (via la matrice d'attention), ce qui le rend plus performant.

Mamba vise à offrir le meilleur des deux mondes. Un petit État aussi puissant qu'un État transformateur :

Comme nous l'avons mentionné plus haut, il le fait en comprimant sélectivement les données dans des états. Lorsque vous avez une phrase en entrée, il y a généralement des informations qui n'ont pas beaucoup de sens, comme un mot d'arrêt.

Afin de compresser sélectivement l'information, nous avons besoin que les paramètres dépendent de l'entrée. À cette fin, explorons d'abord les dimensions des entrées et des sorties du SSM pendant la formation :

Dans le modèle d'espace d'état structuré (S4), la matrice A,B répondre en chantant C sont indépendantes des entrées car leurs dimensions N répondre en chantant D est statique et ne change pas.

Au lieu de cela, Mamba, en combinant la longueur de la séquence d'entrée et la taille du lot, fait de la matrice B répondre en chantant Cmême stimulateur cardiaque ∆_, en fonction de l'apport :

Cela signifie que pour chaque jeton d'entrée, nous avons maintenant différents B répondre en chantant C qui résout le problème de la perception du contenu !

prendre note: Matrix A reste inchangé parce que nous voulons que l'état lui-même reste statique, mais la façon dont il est affecté (par l'intermédiaire de la fonction B répondre en chantant C) est dynamique.

ensemble de manière sélective Sélectionnez les éléments à conserver dans l'état caché et ceux à ignorer, car ils dépendent désormais de l'entrée.

plus petit stimulateur cardiaque ∆ conduit à ignorer des mots spécifiques au profit du contexte antérieur, tandis que de plus grandes stimulateur cardiaque ∆ L'accent est mis davantage sur le vocabulaire d'entrée que sur le contexte :

opération de balayage

Puisque ces matrices sont maintenant dynamique (science) ils ne peuvent pas être calculés à l'aide de la représentation convolutionnelle, qui suppose une fixation du noyau de convolution. Nous ne pouvons utiliser que des représentations récursives, ce qui fait perdre l'avantage de parallélisme offert par la convolution.

Pour parvenir à la parallélisation, examinons comment nous pouvons utiliser la récursivité pour calculer la sortie :

Chaque état est l'état précédent (multiplié par A) avec l'entrée actuelle (multipliée par B) de et. C'est ce qu'on appelle la opération de balayageCette valeur peut être facilement calculée à l'aide d'une boucle "for".

En revanche, la parallélisation semble impossible, car chaque état ne peut être calculé que lorsqu'un état précédent est disponible. Mamba a cependant permis de paralléliser l'état de chaque état en balayage parallèle Les algorithmes rendent cela possible.

Elle suppose que l'ordre dans lequel nous effectuons les opérations n'a pas d'importance, en utilisant une propriété de la loi d'union. Ainsi, nous pouvons calculer la séquence en plusieurs parties, puis les combiner de manière itérative :

matrice dynamique B répondre en chantant C et les algorithmes de balayage parallèle se rejoignent pour créer Algorithme de balayage sélectifpour représenter la nature dynamique et rapide de l'utilisation de représentations récursives.

Algorithmes tenant compte du matériel

L'un des inconvénients des GPU récents est la vitesse de transfert limitée (IO) entre leur SRAM, petite mais efficace, et leur DRAM, grande mais légèrement moins efficace. La copie fréquente d'informations entre la SRAM et la DRAM peut devenir un goulot d'étranglement.

Mamba est similaire à Flash Attention en ce sens qu'il tente de limiter le nombre d'allers-retours entre la DRAM et la SRAM. Pour ce faire, il fusion des noyaux La mise en œuvre de cette procédure permet au modèle d'éviter d'écrire des résultats intermédiaires et de continuer à effectuer le calcul jusqu'à ce qu'il soit terminé.

Nous pouvons voir des exemples spécifiques d'allocation de DRAM et de SRAM en visualisant l'architecture de base de Mamba :

Ici, les éléments suivants sont fusionnés en un seul noyau :

• L'étape de discrétisation est la même que taille du pas ∆
• Algorithme de balayage sélectif
• avec C la multiplication

La dernière partie de l'algorithme tenant compte du matériel est la suivante recalcul.

Les états intermédiaires ne sont pas sauvegardés, mais sont nécessaires pour calculer le gradient dans la passe inverse. Les auteurs recalculent ces états intermédiaires lors de la passe inverse.

Bien que cela puisse sembler inefficace, c'est beaucoup moins coûteux que de lire tous ces états intermédiaires à partir de la DRAM, qui est relativement lente.

Nous avons maintenant couvert toutes les composantes de son architecture, qui est illustrée ci-dessous dans son article :

SSM sélectif. Tiré de : Gu, Albert et Tri Dao, "Mamba : linear time series modelling with selective state spaces". arXiv preprint arXiv:2312.00752 (2023).

Cette architecture est souvent appelée SSM sélectif peut-être S6 car il s'agit essentiellement d'un modèle S4 calculé à l'aide de l'algorithme de balayage sélectif.

Module Mamba

Ce que nous avons exploré jusqu'à présent SSM sélectif peut être mis en œuvre sous la forme d'un module, tout comme nous pouvons représenter l'auto-attention dans un module de décodage.

Comme un décodeur, nous pouvons empiler plusieurs modules Mamba et utiliser leur sortie comme entrée pour le module Mamba suivant :

Il commence par une projection linéaire pour étendre l'intégration de l'entrée. Ensuite, la SSM sélectif La convolution est appliquée au préalable pour empêcher le calcul indépendant des jetons.

SSM sélectif présente les caractéristiques suivantes :

• faire passer (un projet de loi, une inspection, etc.) discret créé SSM récursif
• matrice A aller de l'avant HiPPO Initialisation de la capture dépendance à long terme
• Algorithme de balayage sélectif pour comprimer sélectivement les informations
• Algorithmes tenant compte du matériel pour accélérer les calculs

Lorsque nous examinons la mise en œuvre du code, nous pouvons étendre cette architecture encore davantage et explorer ce à quoi ressemblerait un exemple de bout en bout :

Notez certains changements, tels que l'ajout d'une couche de normalisation et d'un softmax pour la sélection des tokens de sortie.

Lorsque nous mettons tout cela ensemble, nous obtenons une inférence et une formation rapides, même avec un contexte infini. En utilisant cette architecture, les auteurs ont constaté que ses performances égalaient et parfois dépassaient celles d'un modèle Transformer de même taille !

rendre un verdict

Ceci conclut notre exploration des modèles d'espace d'état et de l'incroyable architecture Mamba utilisant des modèles d'espace d'état sélectifs. Nous espérons que cet article vous a permis de mieux comprendre le modèle d'espace d'état et Mamba en particulier. Qui sait si cela remplacera Transformers, mais pour l'instant, il est étonnant de voir une architecture aussi différente recevoir l'attention qu'elle mérite !

Pour voir d'autres visualisations liées à la modélisation du langage et pour soutenir cette lettre d'information, consultez le livre que j'ai co-écrit avec Jay Alammar.