Mamba : Redéfinir la modélisation de séquences et surpasser l’architecture des Transformers

Dans cet article sur Mamba, nous allons explorer comment ce modèle d’espace d’état innovant (SSM) 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 assure des performances exceptionnelles dans ces modalités diverses.

Nous allons nous plonger dans la capacité de Mamba à surmonter les défis computationnels auxquels sont confrontés les Transformers traditionnels, en particulier avec les longues séquences. 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 de la séquence, améliorant ainsi considérablement le débit.

L’originalité de Mamba réside dans sa capacité de traitement rapide, sa couche SSM sélective et sa conception orientée matériel inspirée de FlashAttention. Ces fonctionnalités permettent à Mamba de surpasser de nombreux modèles existants, y compris ceux basés sur l’approche des Transformers, ce qui en fait un progrès notable dans le domaine de l’apprentissage automatique.

Transformers vs Mamba

Les Transformers, 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 prend les devants, avec sa capacité à traiter les longues séquences de manière plus efficace et son architecture unique qui simplifie l’ensemble du processus.

Les Transformers sont habiles dans la gestion de séquences de données, telles que le 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 l’ensemble des séquences simultanément, leur permettant 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 de la prise de décisions.

Cette attention est calculée à l’aide de trois ensembles de poids : requêtes, clés et valeurs, dérivés des données d’entrée. Chaque élément d’une séquence est comparé à chaque autre élément, fournissant un poids qui signale l’importance, ou « attention », que chaque élément devrait recevoir lors de la prédiction du prochain élément de la séquence.

Les Transformers maintiennent 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ête et un réseau feed-forward entièrement connecté simple et positionnel. La normalisation et les connexions résiduelles sont utilisées à chaque sous-couche pour aider à la formation de réseaux profonds.

Le décodeur possède également des couches avec deux sous-couches similaires à celles de l’encodeur, mais ajoute une troisième sous-couche qui effectue une attention multi-tête sur la sortie de l’encodeur. La nature séquentielle du décodeur garantit que les prédictions pour une position ne peuvent considérer que les positions précédentes, préservant la propriété autorégressive.

Contrairement aux Transformers, le modèle Mamba adopte une approche différente. Alors que les Transformers traitent le problème des longues séquences en utilisant des mécanismes d’attention plus complexes, Mamba utilise des espaces d’état sélectifs, offrant une approche plus efficace pour la modélisation de séquences.

Mamba se distingue par sa capacité de traitement rapide, sa couche SSM sélective et sa conception orientée matériel inspirée de FlashAttention. Ces fonctionnalités permettent à Mamba de surpasser de nombreux modèles existants, y compris ceux basés sur l’approche des Transformers, ce qui en fait un progrès notable dans le domaine de l’apprentissage automatique.

Transformers vs Mamba

Les Transformers, 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 prend les devants, avec sa capacité à traiter les longues séquences de manière plus efficace et son architecture unique qui simplifie l’ensemble du processus.

Ce qui rend Mamba vraiment unique, c’est son écart par rapport aux blocs d’attention et de MLP traditionnels. Cette simplification conduit à un modèle plus léger et plus rapide qui se met à l’échelle linéaire avec la longueur de la séquence, une prouesse sans précédent pour ses prédécesseurs.

Les caractéristiques clés de Mamba incluent :

1. SSM 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 une raisonnement basé sur le contenu efficace.
2. Algorithme orienté matériel : Mamba utilise un algorithme parallèle qui est optimisé pour le matériel moderne, en particulier les GPU. Cette conception permet un calcul plus rapide et réduit les exigences de mémoire par rapport aux modèles traditionnels.
3. Architecture simplifiée : En intégrant des SSM sélectifs et en éliminant les blocs d’attention et de MLP, Mamba offre une structure plus simple et plus homogène. Cela conduit à une meilleure mise à l’échelle et à de meilleures performances.

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

Le code et les modèles pré-entraînés de Mamba sont ouverts à la communauté pour une utilisation sur GitHub.

Standard Copying tasks are simple for linear models. Selective Copying and Induction Heads require dynamic, content-aware memory for LLMs.

Les modèles d’espace d’état structurés (S4) sont récemment apparus comme une classe prometteuse de modèles de séquences, englobant les caractéristiques des RNN, des CNN et des modèles d’espace d’état classiques. Les modèles S4 tirent leur inspiration des systèmes continus, en particulier d’un type de système qui mappe des fonctions ou des séquences unidimensionnelles à travers un état latent implicite. Dans le contexte de l’apprentissage automatique, ils représentent une innovation significative, offrant 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 prend une séquence x et produit une sortie y en utilisant des paramètres appris A, B, C et un paramètre de retard Δ. La transformation implique la discrétisation des paramètres (en transformant les fonctions continues en 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 étapes temporelles.

L’importance de la discrétisation

La discrétisation est un processus clé qui transforme les paramètres continus en paramètres discrets à l’aide de 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 la stabilité et les performances du modèle. La discrétisation établit également des parallèles avec les mécanismes de contrôle des flux d’information trouvés dans les RNN, qui sont essentiels pour gérer le flux d’information à travers le réseau.

Linéarité dans le temps (LTI)

Une fonctionnalité clé des modèles S4 est leur linéarité dans le temps. Cette propriété implique que la dynamique du modèle reste cohérente dans le temps, avec les paramètres fixés pour toutes les étapes temporelles. La LTI est une pierre angulaire de la récurrence et des convolutions, offrant un cadre simplifié mais puissant pour construire des modèles de séquences.

Surmonter les limitations fondamentales

Le cadre S4 a été traditionnellement limité par sa nature LTI, qui pose des défis pour modéliser des données qui nécessitent des dynamiques adaptatives. La recherche récente présente une approche qui surmonte ces limitations en introduisant des paramètres à variation dans le temps, supprimant ainsi la contrainte de LTI. Cela permet aux modèles S4 de gérer un ensemble plus diversifié de séquences et de tâches, élargissant considérablement leur applicabilité.

Le terme « modèle d’espace d’état » couvre en général tout processus récurrent impliquant un état latent et a été utilisé pour décrire divers concepts dans plusieurs disciplines. Dans le contexte de l’apprentissage automatique, les modèles S4, ou les SSM structurés, font référence à une classe spécifique de modèles qui ont été optimisés pour une computation efficace 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 de neurones à partir de bout en bout, fonctionnant comme des transformations de séquences autonomes. Ils peuvent être considérés comme analogues aux couches de convolution dans les CNN, fournissant la colonne vertébrale de la modélisation de séquences dans une variété d’architectures de réseaux de neurones.

SSM vs SSM + Selection

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

Structured SSMs

L’article soutient qu’un aspect fondamental de la modélisation de séquences est la compression du contexte en un état gérable. Les modèles qui peuvent sélectivement se concentrer sur ou filtrer les entrées fournissent un moyen plus efficace de maintenir cet état compressé, conduisant à des modèles de séquences plus efficaces et plus puissants. Cette sélectivité est vitale pour que les modèles puissent contrôler de manière adaptative la façon dont l’information s’écoule le long de la dimension de la 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 d’être dépendants de l’entrée, introduisant un degré d’adaptabilité précédemment inatteignable avec les modèles à variation dans le temps. Cela donne lieu à des SSM à variation dans le temps qui ne peuvent plus utiliser des convolutions pour une computation efficace mais qui reposent plutôt sur un mécanisme de récurrence linéaire, une déviation significative des modèles traditionnels.

SSM + Sélection (S6) Cette variante inclut 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 à variation dans le temps en utilisant une opération de balayage, qui traite les éléments de manière séquentielle, ajustant dynamiquement le focus dans le temps.

Performances de Mamba

En termes de performances, Mamba excelle à la fois en vitesse d’inférence et en précision. Sa conception permet une meilleure utilisation de contextes plus longs, ce qui est démontré à la fois dans la modélisation de l’ADN et de l’audio, surpassant les modèles précédents sur des tâches complexes nécessitant des dépendances à longue portée. Sa polyvalence est également mise en évidence dans les évaluations à zéro tir sur plusieurs tâches, établissant une nouvelle norme pour de tels modèles en termes d’efficacité et de scalabilité.

Commencer avec Mamba

Pour ceux qui souhaitent exploiter Mamba, les exigences techniques incluent un système d’exploitation Linux, une carte graphique NVIDIA, PyTorch 1.12+ et CUDA 11.6+. L’installation implique des commandes pip simples pour installer les packages nécessaires à partir du référentiel Mamba. Si des problèmes de compatibilité surviennent avec les versions de PyTorch, utiliser le drapeau –no-build-isolation avec pip peut aider. Ces modèles, formés sur des ensembles de données étendus comme le Pile et le SlimPajama dataset, sont conçus pour répondre à divers besoins computationnels et de performances.

Mamba offre différents niveaux d’interfaces, allant de la couche SSM sélective au bloc Mamba et aux structures de modèles de langage complètes. Le bloc Mamba, qui est le module principal de l’architecture, utilise une couche Conv1d causale et peut être facilement intégré dans les conceptions de réseaux de neurones. L’exemple d’utilisation en Python démontre la création d’un modèle Mamba et le traitement de données à travers celui-ci, mettant en évidence la simplicité et la flexibilité du système.

Les modèles pré-entraînés de Mamba sont disponibles sur Hugging Face, avec des tailles allant de 130M à 2,8G de paramètres, formés sur l’ensemble de données étendu Pile et le SlimPajama dataset. Ces modèles sont conçus pour répondre à des besoins computationnels et de performances divers, en adhérant aux normes dimensionnelles du GPT-3. Les utilisateurs peuvent s’attendre à un débit élevé et à une grande précision avec ces modèles, ce qui fait de Mamba un choix compétitif pour diverses applications, y compris mais sans s’y limiter, la modélisation du langage.

L’impact de Mamba

Mamba représente un bond en avant dans la modélisation de séquences, offrant une alternative puissante aux architectures de Transformers pour le traitement de données denses en information. Sa conception s’aligne sur les exigences du matériel moderne, en optimisant à la fois l’utilisation de la mémoire et les capacités de traitement parallèle. La disponibilité en open source du codebase 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 dans le domaine de l’IA et de l’apprentissage automatique.