Intelligence artificielle
Mamba : Révolutionner la modélisation de séquences et surpasser l’architecture des Transformers
Dans cet article sur Mamba, nous allons explorer comment ce modèle innovant d’espace d’état (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 linéaires avec des espaces d’état sélectifs assure des performances exceptionnelles dans ces modalités diverses.
Nous allons examiner la capacité de Mamba à surmonter les défis computationnels auxquels sont confrontés les Transformers traditionnels, en particulier avec des 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 de la séquence, 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 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 une avancée notable dans l’apprentissage automatique.
Transformers vs Mamba
Les Transformers, comme GPT-4, ont établi des références en traitement du langage naturel. Cependant, leur efficacité diminue avec des séquences plus longues. C’est là que Mamba prend le dessus, avec sa capacité à traiter des séquences longues de manière plus efficace et son architecture unique qui simplifie l’ensemble du processus.
Les Transformers sont aptes à traiter 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 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 prédiction.
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 de l’élément suivant 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 fully connected simple et position-wise. La normalisation et les connexions résiduelles sont utilisées à chaque sous-couche pour aider à la formation de réseaux profonds.
Le décodeur a é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é auto-régressive.
Contrairement aux Transformers, le modèle Mamba adopte une approche différente. Alors que les Transformers traitent le problème des séquences longues en utilisant des mécanismes d’attention plus complexes, Mamba utilise des espaces d’état sélectifs, offrant une approche plus efficace et plus simple pour traiter les séquences.
Voici une vue d’ensemble de haut niveau de la façon dont fonctionne un Transformer :
- Traitement des entrées : Les Transformers codent d’abord les données d’entrée dans un format que le modèle peut comprendre, souvent en utilisant des embeddings qui incorporrent également la position de chaque élément de la séquence.
- Mécanisme d’attention : Au cœur du mécanisme d’attention, il calcule un score qui représente à quel point il faut se concentrer sur d’autres parties de la séquence d’entrée lors de la compréhension d’un élément actuel.
- Architecture Encodeur-Décodeur : Le modèle Transformer est composé d’un encodeur pour traiter les entrées et d’un décodeur pour générer la sortie. Chacun se compose de plusieurs couches qui raffinent la compréhension du modèle des entrées.
- Attention multi-tête : À la fois dans l’encodeur et le décodeur, l’attention multi-tête permet au modèle de se concentrer simultanément sur différentes parties de la séquence à partir de différents espaces de représentation, améliorant ainsi sa capacité à apprendre de contextes divers.
- Réseaux feed-forward position-wise : Après l’attention, un réseau neural simple traite la sortie de chaque position séparément et de manière identique. Cela est combiné avec l’entrée à travers une connexion résiduelle et suivi d’une normalisation de couche.
- 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 par ce qu’il a généré jusqu’à présent.
La capacité du Transformer à traiter 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.
Contrairement à cela, le modèle Mamba fonctionne différemment en utilisant des espaces d’état sélectifs pour traiter les séquences. Cette approche répond à l’inefficacité computationnelle des Transformers lorsqu’ils traitent des séquences longues. La conception de Mamba permet une inférence plus rapide et une mise à l’échelle linéaire avec la longueur de la séquence, établissant un nouveau paradigme pour la modélisation de séquences qui pourrait être plus efficace, en particulier à mesure que les séquences deviennent de plus en plus longues.
Mamba
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éairement avec la longueur de la séquence – une prouesse sans égale par ses prédécesseurs.
Les fonctionnalités clés de Mamba incluent :
- 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.
- Algorithme conscient du matériel : Mamba utilise un algorithme parallèle 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.
- 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 à des performances améliorées.
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é-formation et de formation spécifiques au domaine. Par exemple, dans la modélisation du langage, Mamba égale ou dépasse les performances de modèles Transformer plus importants.
Le code et les modèles pré-formés de Mamba sont disponibles en open source pour la communauté sur GitHub.
Les tâches de copie standard sont simples pour les modèles linéaires. La copie sélective et les têtes d’induction nécessitent une mémoire dynamique et sensible au contenu pour les LLM.
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 des 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 unidimensionnelles ou des séquences à travers un état latent implicite. Dans le contexte de l’apprentissage profond, ils représentent une innovation significative, offrant une nouvelle méthodologie pour concevoir des modèles de séquences qui sont 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 convertissant des 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 ainsi aux modèles S4 de maintenir une connexion avec les systèmes à temps continu. Cela dote les modèles de 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 verrouillage trouvés dans les RNN, qui sont critiques pour gérer le flux d’informations à travers le réseau.
Invariance linéaire dans le temps (LTI)
Une caractéristique fondamentale des modèles S4 est leur invariance linéaire 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. LTI est un élément essentiel 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 dans la modélisation de données nécessitant des dynamiques adaptatives. L’article de recherche 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 profond, les modèles S4, ou les modèles d’espace d’état 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 le squelette pour la modélisation de séquences dans diverses architectures de réseaux de neurones.
SSM vs SSM + Sélection
Motivation pour la sélectivité dans la modélisation de séquences
Modèles d’espace d’état structurés
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 offrent un moyen plus efficace de maintenir cet état compressé, conduisant ainsi à des modèles de séquences 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 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 des entrées, introduisant ainsi un degré d’adaptabilité qui n’était pas atteignable avec les modèles à temps invariable. Cela donne lieu à des SSM à variation dans le temps qui ne peuvent plus utiliser des convolutions pour une computation efficace mais qui s’appuient plutôt sur un mécanisme de récurrence linéaire, une déviation significative par rapport aux modèles traditionnels.
SSM + Sélection (S6) Cette variante inclut un mécanisme de sélection, ajoutant une dépendance aux entrées pour les 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 séquentiellement, ajustant dynamiquement le focus au fil du temps.
Performances de Mamba
Mamba est le meilleur de sa catégorie pour chaque résultat d’évaluation
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 dans 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 pour plusieurs tâches, établissant un nouveau standard pour de tels modèles en termes d’efficacité et de mise à l’échelle.
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 nécessite de simples commandes pip 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 aux références 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 des conceptions de réseaux de neurones. L’exemple d’utilisation en Python démontre la simplicité et la flexibilité du système en instanciant un modèle Mamba et en traitant des données à travers celui-ci.
Les modèles pré-formés de Mamba sont disponibles sur Hugging Face, avec des tailles allant de 130M à 2,8B de paramètres, formés sur l’ensemble de données étendu du Pile et du SlimPajama dataset. Ces modèles sont conçus pour répondre à des besoins computationnels et de performances divers, en respectant les 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 non limité à la modélisation du langage.
Impact de Mamba
Mamba représente un bond en avant dans la modélisation de séquences, offrant une alternative puissante aux architectures des Transformers pour traiter des 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é-formé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 profond.
