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

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

December 1, 2023

Auteurs: Albert Gu, Tri Dao

cs.AI

Résumé

Les modèles de fondation, qui alimentent aujourd'hui la plupart des applications passionnantes en apprentissage profond, reposent presque universellement sur l'architecture Transformer et son module d'attention central. De nombreuses architectures à temps sous-quadratique, telles que l'attention linéaire, les convolutions à porte, les modèles récurrents et les modèles à espace d'état structuré (SSM), ont été développées pour remédier à l'inefficacité computationnelle des Transformers sur les longues séquences, mais elles n'ont pas égalé les performances de l'attention sur des modalités importantes comme le langage. Nous identifions qu'une faiblesse clé de ces modèles est leur incapacité à effectuer un raisonnement basé sur le contenu, et nous apportons plusieurs améliorations. Premièrement, simplement permettre aux paramètres du SSM d'être des fonctions de l'entrée résout leur faiblesse avec les modalités discrètes, permettant au modèle de propager ou d'oublier sélectivement l'information le long de la dimension de la longueur de la séquence en fonction du token actuel. Deuxièmement, bien que ce changement empêche l'utilisation de convolutions efficaces, nous concevons un algorithme parallèle optimisé pour le matériel en mode récurrent. Nous intégrons ces SSM sélectifs dans une architecture de réseau neuronal simplifiée de bout en bout sans attention ni même blocs MLP (Mamba). Mamba bénéficie d'une inférence rapide (un débit 5 fois supérieur à celui des Transformers) et d'une mise à l'échelle linéaire avec la longueur de la séquence, et ses performances s'améliorent sur des données réelles jusqu'à des séquences de longueur million. En tant que modèle de séquence général, Mamba atteint des performances de pointe sur plusieurs modalités telles que le langage, l'audio et la génomique. En modélisation du langage, notre modèle Mamba-3B surpasse les Transformers de même taille et égale les Transformers deux fois plus grands, à la fois en pré-entraînement et en évaluation en aval.

English

Foundation models, now powering most of the exciting applications in deep learning, are almost universally based on the Transformer architecture and its core attention module. Many subquadratic-time architectures such as linear attention, gated convolution and recurrent models, and structured state space models (SSMs) have been developed to address Transformers' computational inefficiency on long sequences, but they have not performed as well as attention on important modalities such as language. We identify that a key weakness of such models is their inability to perform content-based reasoning, and make several improvements. First, simply letting the SSM parameters be functions of the input addresses their weakness with discrete modalities, allowing the model to selectively propagate or forget information along the sequence length dimension depending on the current token. Second, even though this change prevents the use of efficient convolutions, we design a hardware-aware parallel algorithm in recurrent mode. We integrate these selective SSMs into a simplified end-to-end neural network architecture without attention or even MLP blocks (Mamba). Mamba enjoys fast inference (5times higher throughput than Transformers) and linear scaling in sequence length, and its performance improves on real data up to million-length sequences. As a general sequence model backbone, Mamba achieves state-of-the-art performance across several modalities such as language, audio, and genomics. On language modeling, our Mamba-3B model outperforms Transformers of the same size and matches Transformers twice its size, both in pretraining and downstream evaluation.