Suivez nous sur

Mamba : redĂ©finir la modĂ©lisation des sĂ©quences et surpasser l'architecture des transformateurs

Intelligence Artificielle

Mamba : redĂ©finir la modĂ©lisation des sĂ©quences et surpasser l'architecture des transformateurs

mm
Modèle d'IA Mamba

Dans cet article sur MambaNous explorerons comment ce modèle d'espace d'état (SSM) innovant révolutionne la modélisation de séquences. Développé par Albert Gu et Tri Dao, Mamba se distingue par son efficacité dans le traitement de séquences complexes dans des domaines tels que le traitement du langage, la génomique et l'analyse audio. Sa modélisation de séquences en temps linéaire avec des espaces d'état sélectifs garantit des performances exceptionnelles dans ces diverses modalités.

Nous explorerons la capacité de Mamba à surmonter les défis de calcul rencontrés par les Transformers traditionnels, notamment avec les séquences longues. Son approche sélective dans les modèles d'espace d'état permet une inférence plus rapide et une mise à l'échelle linéaire avec la longueur des séquences, améliorant ainsi considérablement le débit.

La particularité de Mamba réside dans sa capacité de traitement rapide, sa couche SSM sélective et sa conception matérielle optimisée, inspirée de FlashAttention. Ces fonctionnalités permettent à Mamba de surpasser de nombreux modèles existants, notamment ceux basés sur l'approche Transformer, ce qui en fait une avancée majeure dans l'apprentissage automatique.

Transformateurs vs Mamba 

Les transformateurs, comme GPT-4, ont Ă©tabli des rĂ©fĂ©rences en matière de traitement du langage naturel. Cependant, leur efficacitĂ© diminue avec les sĂ©quences plus longues. C'est lĂ  que Mamba fait un bond en avant : sa capacitĂ© Ă  traiter les sĂ©quences longues plus efficacement et son architecture unique qui simplifie l'ensemble du processus.

Transformateurs capables de gérer des séquences de données, telles que du texte pour les modèles de langage. Contrairement aux modèles précédents qui traitaient les données de manière séquentielle, les Transformers traitent des séquences entières simultanément, ce qui leur permet de capturer des relations complexes au sein des données.

Ils utilisent un mécanisme d'attention, qui permet au modèle de se concentrer sur différentes parties de la séquence lors des prédictions.

Cette attention est calculĂ©e Ă  l'aide de trois ensembles de pondĂ©rations : requĂŞtes, clĂ©s et valeurs, dĂ©rivĂ©s des donnĂ©es d'entrĂ©e. Chaque Ă©lĂ©ment d'une sĂ©quence est comparĂ© Ă  tous les autres, ce qui fournit une pondĂ©ration indiquant l'importance, ou « attention Â», que chaque Ă©lĂ©ment doit recevoir lors de la prĂ©diction de l'Ă©lĂ©ment suivant de la sĂ©quence.

Les transformateurs conservent deux blocs principaux : l'encodeur, qui traite les donnĂ©es d'entrĂ©e, et le dĂ©codeur, qui gĂ©nère la sortie. L'encodeur se compose de plusieurs couches, chacune contenant deux sous-couches : un mĂ©canisme d'auto-attention multi-tĂŞtes et un rĂ©seau de rĂ©troaction simple, entièrement connectĂ© en termes de position. La normalisation et les connexions rĂ©siduelles sont utilisĂ©es Ă  chaque sous-couche pour aider Ă  la formation des rĂ©seaux profonds.

Le décodeur possède également deux sous-couches similaires à celles de l'encodeur, mais ajoute une troisième sous-couche qui effectue une attention multi-têtes sur la sortie de l'encodeur. La nature séquentielle du décodeur garantit que les prédictions pour une position donnée ne prennent en compte que les positions antérieures, préservant ainsi la propriété autorégressive.

Contrairement à Transformers, le modèle Mamba adopte une approche différente. Alors que Transformers traite le problème des longues séquences en utilisant des mécanismes d'attention plus complexes, Mamba utilise des espaces d'états sélectifs, offrant une approche plus informatique.

Voici un aperçu de haut niveau du fonctionnement d’un transformateur :

  1. Traitement des entrées: Les transformateurs codent d'abord les données d'entrée dans un format que le modèle peut comprendre, souvent en utilisant des intégrations qui intègrent également la position de chaque élément dans la séquence.
  2. Mécanisme d'attention: À la base, le mécanisme d'attention calcule un score qui représente le niveau d'attention à accorder aux autres parties de la séquence d'entrée lors de la compréhension d'un élément actuel.
  3. Architecture codeur-décodeur: Le modèle de transformateur est composé d'un encodeur pour traiter l'entrée et d'un décodeur pour générer la sortie. Chacun est constitué de plusieurs couches qui affinent la compréhension de l'entrée par le modèle.
  4. Attention multi-têtes: Au sein de l'encodeur et du décodeur, l'attention multi-têtes permet au modèle de s'occuper simultanément de différentes parties de la séquence à partir de différents espaces de représentation, améliorant ainsi sa capacité à apprendre de divers contextes.
  5. Réseaux Feed-Forward par position: Après attention, un simple réseau de neurones traite la sortie de chaque position séparément et de manière identique. Ceci est combiné avec l’entrée via une connexion résiduelle et suivi d’une normalisation de couche.
  6. Génération de sortie:Le décodeur prédit ensuite une séquence de sortie, influencée par le contexte de l'encodeur et ce qu'il a généré jusqu'à présent.

La capacité du transformateur à gérer des séquences en parallèle et son mécanisme d'attention robuste le rendent puissant pour des tâches telles que la traduction et la génération de texte.

En revanche, le modèle Mamba fonctionne différemment en utilisant des espaces d'état sélectifs pour traiter les séquences. Cette approche corrige l'inefficacité de calcul des Transformers face aux séquences longues. La conception de Mamba permet une inférence plus rapide et une évolutivité linéaire avec la longueur des séquences, établissant un nouveau paradigme pour la modélisation des séquences, qui pourrait être plus efficace, notamment lorsque les séquences deviennent de plus en plus longues.

Mamba

Ce qui rend Mamba vraiment unique, c'est sa rupture avec l'attention traditionnelle et les blocages MLP. Cette simplification conduit à un modèle plus léger et plus rapide qui évolue linéairement avec la longueur de la séquence – un exploit inégalé par ses prédécesseurs.

Les principales fonctionnalitĂ©s de Mamba incluent :

  1. MSU sélectifs: Ceux-ci permettent à Mamba de filtrer les informations non pertinentes et de se concentrer sur les données pertinentes, améliorant ainsi sa gestion des séquences. Cette sélectivité est cruciale pour un raisonnement efficace basé sur le contenu.
  2. Algorithme sensible au matérielMamba utilise un algorithme parallèle optimisé pour le matériel moderne, notamment les GPU. Cette conception permet des calculs plus rapides et réduit les besoins en mémoire par rapport aux modèles traditionnels.
  3. Architecture simplifiée: En intégrant des SSM sélectifs et en éliminant les blocages d'attention et MLP, Mamba propose une structure plus simple et plus homogène. Cela conduit à une meilleure évolutivité et de meilleures performances.

Mamba a démontré des performances supérieures dans divers domaines, notamment le langage, l'audio et la génomique, excellant à la fois dans les tâches de pré-formation et spécifiques à un domaine. Par exemple, dans la modélisation du langage, Mamba égale ou dépasse les performances des modèles Transformer plus grands.

Le code et les modèles pré-entraînés de Mamba sont ouvertement disponibles pour une utilisation communautaire à l'adresse GitHub.

Les tâches de copie standard sont simples pour les modèles linéaires. Les têtes de copie sélective et d'induction nécessitent une mémoire dynamique et sensible au contenu pour les LLM.

Les tâches de copie standard sont simples pour les modèles linéaires. Les têtes de copie sélective et d'induction nécessitent une mémoire dynamique et sensible au contenu pour les LLM.

Les modèles d'espace d'état structuré (S4) sont récemment apparus comme une classe prometteuse de modèles de séquence, englobant les caractéristiques des RNN, des CNN et des modèles d'espace d'état classiques. Les modèles S4 s'inspirent des systèmes continus, en particulier d'un type de système qui mappe des fonctions ou des séquences unidimensionnelles via un état latent implicite. Dans le contexte de l’apprentissage profond, ils représentent une innovation significative, fournissant une nouvelle méthodologie pour concevoir des modèles de séquences efficaces et hautement adaptables.

La dynamique des modèles S4

SSM (S4) Il s’agit du modèle d’espace d’état structuré de base. Il faut une séquence x et produit une sortie y en utilisant les paramètres appris A, B, C, et un paramètre de retard ΔLa transformation implique la discrétisation des paramètres (transformation des fonctions continues en fonctions discrètes) et l'application de l'opération SSM, qui est invariante dans le temps, ce qui signifie qu'elle ne change pas au fil des différentes étapes de temps.

L'importance de la discrétisation

La discrétisation est un processus clé qui transforme les paramètres continus en paramètres discrets via des formules fixes, permettant aux modèles S4 de maintenir une connexion avec les systèmes à temps continu. Cela confère aux modèles des propriétés supplémentaires, telles que l'invariance de résolution, et garantit une normalisation appropriée, améliorant ainsi la stabilité et les performances du modèle. La discrétisation établit également des parallèles avec les mécanismes de contrôle présents dans les RNN, qui sont essentiels à la gestion du flux d'informations à travers le réseau.

Invariance temporelle linéaire (LTI)

Une caractéristique essentielle des modèles S4 est leur invariance temporelle linéaire. Cette propriété implique que la dynamique du modèle reste cohérente dans le temps, avec les paramètres fixes pour tous les pas de temps. LTI est la pierre angulaire de la récurrence et des convolutions, offrant un cadre simplifié mais puissant pour créer des modèles de séquence.

Surmonter les limites fondamentales

Le cadre S4 a été traditionnellement limité par sa nature LTI, ce qui pose des défis dans la modélisation de données nécessitant une dynamique adaptative. Le récent document de recherche présente une approche qui surmonte ces limitations en introduisant des paramètres variables dans le temps, supprimant ainsi la contrainte du LTI. Cela permet aux modèles S4 de gérer un ensemble plus diversifié de séquences et de tâches, élargissant ainsi considérablement leur applicabilité.

Le terme « modèle d'espace d'Ă©tat Â» dĂ©signe au sens large tout processus rĂ©current impliquant un Ă©tat latent et a Ă©tĂ© utilisĂ© pour dĂ©crire divers concepts dans de nombreuses disciplines. Dans le contexte de l'apprentissage profond, les modèles S4, ou SSM structurĂ©s, dĂ©signent une classe spĂ©cifique de modèles optimisĂ©s pour des calculs efficaces tout en conservant la capacitĂ© de modĂ©liser des sĂ©quences complexes.

Les modèles S4 peuvent être intégrés dans des architectures de réseaux neuronaux de bout en bout, fonctionnant comme des transformations de séquence autonomes. Ils peuvent être considérés comme analogues aux couches de convolution des CNN, fournissant l’épine dorsale de la modélisation de séquences dans diverses architectures de réseaux neuronaux.

SSM vs SSM + Sélection

SSM vs SSM + Sélection

Motivation pour la sélectivité dans la modélisation de séquences

SSM structurés

SSM structurés

L'article soutient qu'un aspect fondamental de la modélisation de séquences est la compression du contexte dans un état gérable. Les modèles qui peuvent se concentrer ou filtrer sélectivement les entrées fournissent un moyen plus efficace de maintenir cet état compressé, conduisant à des modèles de séquence plus efficaces et plus puissants. Cette sélectivité est vitale pour que les modèles contrôlent de manière adaptative la façon dont les informations circulent le long de la dimension de séquence, une capacité essentielle pour gérer des tâches complexes dans la modélisation du langage et au-delà.

Les SSM sélectifs améliorent les SSM conventionnels en permettant à leurs paramètres de dépendre des entrées, ce qui introduit un degré d'adaptabilité auparavant inaccessible avec les modèles invariants dans le temps. Il en résulte des SSM variables dans le temps qui ne peuvent plus utiliser des convolutions pour un calcul efficace, mais s'appuient plutôt sur un mécanisme de récurrence linéaire, un écart significatif par rapport aux modèles traditionnels.

SSM + Sélection (S6) Cette variante comprend un mécanisme de sélection, ajoutant une dépendance à l'entrée aux paramètres B et C, et un paramètre de retard Δ. Cela permet au modèle de se concentrer sélectivement sur certaines parties de la séquence d'entrée. x. Les paramètres sont discrétisés en tenant compte de la sélection, et l'opération SSM est appliquée de manière variable dans le temps à l'aide d'une opération de balayage, qui traite les éléments de manière séquentielle, en ajustant la focalisation de manière dynamique au fil du temps.

Points forts des performances de Mamba

Mamba est le meilleur de sa catégorie pour chaque résultat d'évaluation

Mamba est le meilleur de sa catégorie pour chaque résultat d'évaluation

En termes de performances, Mamba excelle tant en vitesse d'infĂ©rence qu'en prĂ©cision. Sa conception permet une meilleure exploitation des contextes plus longs, comme en tĂ©moignent les modĂ©lisations ADN et audio, surpassant les modèles prĂ©cĂ©dents pour les tâches complexes nĂ©cessitant des dĂ©pendances Ă  longue portĂ©e. Sa polyvalence est Ă©galement mise en Ă©vidence lors des Ă©valuations « zero-shot Â» sur plusieurs tâches, Ă©tablissant ainsi une nouvelle norme pour ces modèles en termes d'efficacitĂ© et d'Ă©volutivitĂ©.

Débuter avec Mamba

Pour ceux qui souhaitent tirer parti de Mamba, les exigences techniques incluent un système d'exploitation Linux, un GPU NVIDIA, PyTorch 1.12+ et CUDA 11.6+. L'installation implique de simples commandes pip pour installer les packages nécessaires à partir du Dépôt Mamba. Si des problèmes de compatibilité surviennent avec les versions de PyTorch, l'utilisation de l'indicateur –no-build-isolation avec pip peut aider. Ces modèles, formés sur des ensembles de données étendus tels que les ensembles de données Pile et SlimPajama, sont conçus pour répondre à divers besoins informatiques et références de performances.

Mamba propose différents niveaux d'interfaces, de la couche SSM sélective au bloc Mamba et aux structures complètes du modèle de langage. Le bloc Mamba, module principal de l'architecture, utilise une couche causale Conv1d et s'intègre facilement aux conceptions de réseaux neuronaux. L'exemple d'utilisation en Python fourni illustre l'instanciation d'un modèle Mamba et le traitement des données via celui-ci, soulignant ainsi la simplicité et la flexibilité du système.

Les modèles Mamba prĂ©-entraĂ®nĂ©s sont disponibles sur Étreindre le visage, avec des tailles allant de 130 M Ă  2.8 B de paramètres, formĂ©s sur le vaste ensemble de donnĂ©es Pile et l'ensemble de donnĂ©es SlimPajama. Ces modèles sont conçus pour rĂ©pondre Ă  diverses exigences de calcul et de performances, en respectant les normes dimensionnelles de GPT-3. Les utilisateurs peuvent s'attendre Ă  un dĂ©bit et une prĂ©cision Ă©levĂ©s de la part de ces modèles, faisant de Mamba un choix compĂ©titif pour diverses applications, y compris, mais sans s'y limiter, la modĂ©lisation de langage.

L'impact de Mamba

Mamba représente une avancée majeure dans la modélisation séquentielle, offrant une alternative performante aux architectures Transformer pour le traitement de données denses en informations. Sa conception répond aux exigences du matériel moderne, optimisant à la fois l'utilisation de la mémoire et les capacités de traitement parallèle. La disponibilité open source de la base de code de Mamba et de ses modèles pré-entraînés en fait un outil accessible et robuste pour les chercheurs et les développeurs en IA et en apprentissage profond.

J'ai passé les cinq dernières années à m'immerger dans le monde fascinant du Machine Learning et du Deep Learning. Ma passion et mon expertise m'ont amené à contribuer à plus de 50 projets de génie logiciel divers, avec un accent particulier sur l'IA/ML. Ma curiosité continue m'a également attiré vers le traitement automatique du langage naturel, un domaine que j'ai hâte d'explorer davantage.