Intelligence artificielle
Mamba : redéfinir la modélisation des séquences et surpasser l'architecture des transformateurs
Dans cet article sur Mamba, nous explorerons comment ce modèle innovant d'espace d'états (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 temporelles linéaires avec des espaces d'états sélectifs garantit des performances exceptionnelles dans ces diverses modalités.
Nous examinerons la capacité de Mamba à surmonter les défis informatiques rencontrés par les Transformers traditionnels, en particulier avec de longues séquences. Son approche sélective dans les modèles d’espace d’états permet une inférence plus rapide et une mise à l’échelle linéaire avec la longueur de séquence, améliorant considérablement le débit.
Le caractère unique de Mamba réside dans sa capacité de traitement rapide, sa couche SSM sélective et sa conception conviviale inspirée de FlashAttention. Ces fonctionnalités permettent à Mamba de surpasser de nombreux modèles existants, y compris ceux basés sur l'approche Transformer, ce qui en fait une avancée remarquable 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 des séquences plus longues. C'est ici que Mamba avance, avec sa capacité à traiter de longues séquences 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ées des données d'entrée. Chaque élément d'une séquence est comparé à tous les autres éléments, fournissant un poids qui signifie l'importance, ou « l'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 comporte également des couches avec deux sous-couches similaires à 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 ne peuvent prendre 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 général du fonctionnement d'un transformateur :
- 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.
- 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.
- 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. Chacune se compose de plusieurs couches qui affinent la compréhension de l'entrée par le modèle.
- 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.
- 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.
- 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’états sélectifs pour traiter les séquences. Cette approche résout l’inefficacité informatique des Transformers lorsqu’il s’agit de longues séquences. La conception de Mamba permet une inférence plus rapide et évolue linéairement avec la longueur de la séquence, établissant ainsi 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 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 :
- 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.
- Algorithme sensible au 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 besoins en mémoire par rapport aux modèles traditionnels.
- 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 librement disponibles pour une utilisation communautaire sur 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 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 de discrétiser les paramètres (en transformant des fonctions continues en fonctions discrètes) et d'appliquer l'opération SSM, qui est invariante dans le temps, ce qui signifie qu'elle ne change pas au fil des différents pas 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 » couvre largement 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 SSM structurés, font référence à une classe spécifique de modèles optimisés pour un calcul 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 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
Motivation pour la sélectivité dans la modélisation de séquences
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 ainsi que le 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
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 ADN et 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 sans tir sur plusieurs tâches, établissant une nouvelle norme pour de tels 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 de modèle de langage. Le bloc Mamba, qui est le module principal de l'architecture, utilise une couche causale Conv1d et peut être facilement intégré aux conceptions de réseaux neuronaux. L'exemple d'utilisation fourni en Python montre l'instanciation d'un modèle Mamba et le traitement des données via celui-ci, soulignant 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 un bond en avant dans la modélisation de séquences, offrant une alternative puissante aux architectures Transformer pour le traitement de données riches en informations. Sa conception s'aligne sur les 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 développeurs dans le domaine de l'IA et de l'apprentissage profond.
