Mamba : Modélisation de séquences en temps linéaire avec des espaces d'état sélectifs

Les modèles fondamentaux, qui alimentent aujourd’hui la majorité des applications passionnantes en apprentissage profond, reposent presque universellement sur l’architecture Transformer et son module central d’attention. De nombreuses architectures à temps presque quadratique, telles que l’attention linéaire, les convolutions à entrées filtrées, les modèles récurrents et les modèles d’état structurés (SSMs), ont été développées pour remédier à l’inefficacité computationnelle des Transformers sur des séquences longues. Toutefois, ces approches n’ont pas atteint les performances de l’attention sur des modalités cruciales comme le langage. Nous identifions ici un défaut fondamental de ces modèles : leur incapacité à réaliser un raisonnement fondé sur le contenu. Nous proposons plusieurs améliorations. Premièrement, en permettant que les paramètres du SSM dépendent des adresses d’entrée, nous surmontons cette faiblesse dans les modalités discrètes, permettant au modèle de propager ou d’oublier sélectivement l’information le long de la dimension de longueur de séquence, en fonction du jeton courant. Deuxièmement, bien que ce changement rende impossible l’utilisation des convolutions efficaces, nous concevons un algorithme parallèle adapté aux ressources matérielles, fonctionnant en mode récurrent. Nous intégrons ces SSMs sélectifs dans une architecture de réseau neuronal simplifiée, sans attention ni blocs MLP (Mamba). Mamba bénéficie d’une inférence rapide (débit 5 fois supérieur à celui des Transformers), d’une mise à l’échelle linéaire en longueur de séquence, et ses performances s’améliorent sur des données réelles jusqu’à des séquences de longueur millésimale. En tant que squelette général pour les modèles de séquences, Mamba atteint des performances de pointe sur plusieurs modalités, notamment le langage, l’audio et la génomique. En modélisation du langage, notre modèle Mamba-3B dépasse les Transformers de même taille et égale les performances des Transformers de taille double, tant en phase de pré-entraînement qu’en évaluation sur tâches spécifiques.