Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu; Tri Dao

Articles · arXiv 2312.00752

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

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu, Tri Dao

Publié le 2023-12-01cs.LGcs.AI arXiv PDF

Lire avec traduction et IA

Résumé en trois lignes

🎯 Les Transformers ont une complexité quadratique selon la longueur de séquence et requièrent de volumineux caches KV en inférence autorégressive, tandis que les modèles sous-quadratiques antérieurs comme les modèles d'espace d'états LTI peinent sur les tâches discrètes de raisonnement sur le contenu comme la modélisation du langage.
🔧 Rend les paramètres des SSM dépendants de l'entrée (sélectifs) pour filtrer ou propager dynamiquement l'information, et associe cela à un scan associatif parallèle fusionné adapté au matériel avec recalcul pour éviter de matérialiser l'état caché étendu dans la mémoire HBM du GPU.
📈 Mamba-3B égale ou surpasse des Transformers deux fois plus volumineux sur des tâches aval en zéro-shot (ex. 63,3 % de moyenne contre 59,1 % pour Pythia-2.8B et 61,7 % pour Pythia-6.9B) tout en offrant un débit d'inférence jusqu'à 5× supérieur.

Résumé

Les modèles de fondation, qui alimentent désormais la plupart des applications majeures en apprentissage profond, reposent presque universellement sur l'architecture Transformer et son module central d'attention. De nombreuses architectures à complexité temporelle sous-quadratique, telles que l'attention linéaire, les modèles récurrents ou à convolutions avec portes, ainsi que les modèles d'espace d'états structurés (SSM), ont été développées pour pallier l'inefficacité computationnelle des Transformers sur de longues séquences ; toutefois, elles n'ont pas atteint les performances de l'attention sur des modalités essentielles comme le langage. Nous identifions qu'une faiblesse déterminante de ces modèles réside dans leur incapacité à effectuer un raisonnement fondé sur le contenu, et nous proposons plusieurs améliorations. Premièrement, le simple fait de définir les paramètres des SSM comme des fonctions de l'entrée résout leur lacune sur les modalités discrètes, permettant ainsi au modèle de propager ou d'oublier sélectivement l'information le long de la séquence en fonction du jeton courant. Deuxièmement, bien que ce changement empêche le recours aux convolutions efficaces, nous concevons un algorithme parallèle adapté au matériel en mode récurrent. Nous intégrons ces SSM sélectifs dans une architecture de réseau neuronal de bout en bout simplifiée, sans attention ni même blocs MLP (Mamba). Mamba bénéficie d'une inférence rapide (débit $5 \times$ supérieur à celui des Transformers) et d'un passage à l'échelle linéaire avec la longueur de séquence, et ses performances continuent de progresser sur des données réelles jusqu'à des séquences d'un million d'éléments. En tant que squelette générique pour la modélisation de séquences, Mamba établit l'état de l'art 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 des Transformers deux fois plus grands, tant en préentraînement qu'en évaluation aval.

Résumé original (anglais)

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 (5 $\times$ 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.

Guide de lecture

Problème et motivation

Les modèles de fondation reposent massivement sur l'architecture Transformer et l'auto-attention. Bien que l'attention dense achemine efficacement l'information à travers une fenêtre de contexte séquentielle, elle souffre de deux limites fondamentales : un coût de calcul quadratique $O (L^{2})$ par rapport à la longueur de séquence $L$ , et la nécessité de stocker les clés et valeurs passées (le cache KV), ce qui entraîne une inférence par pas en temps linéaire et des goulots d'étranglement mémoire.

Les approches sous-quadratiques antérieures — telles que l'attention linéaire, les architectures récurrentes et les modèles d'espace d'états structurés (SSM) comme S4 — évoluent de manière linéaire ou pseudo-linéaire en temps. Cependant, les SSM structurés n'ont historiquement pas réussi à égaler la qualité des Transformers sur des modalités discrètes et denses en information telles que le texte. Les auteurs en identifient la cause principale : les SSM antérieurs sont linéaires et invariants dans le temps (LTI). Leurs dynamiques de transition ne variant pas en fonction du contenu de l'entrée, ils ne peuvent pas compresser sélectivement le contexte, filtrer les bruits non pertinents ou moduler la rétention à travers des espacements variables.

Idées maîtresses

Mécanisme de sélection comme compression de contexte : La modélisation de séquences nécessite de compresser l'historique dans un état. Les SSM LTI antérieurs utilisent des paramètres de transition fixes, ce qui les contraint à traiter tous les pas de temps de manière uniforme. En paramétrant les matrices du SSM en fonction du jeton d'entrée courant, le modèle acquiert la sensibilité au contenu requise pour filtrer ou mémoriser sélectivement les informations.
Surmonter le goulot d'étranglement de la non-convolution : Rendre les paramètres dépendants de l'entrée brise l'équivalence mathématique entre récurrences linéaires et convolutions globales. Sans convolution, une récurrence naïve nécessiterait de matérialiser un état latent étendu de taille $(B, L, D, N)$ dans la mémoire à large bande passante (HBM) du GPU. Mamba résout ce problème en fusionnant la discrétisation, le scan associatif parallèle et la projection de sortie au sein de la SRAM rapide sur puce du GPU, réduisant au minimum le trafic vers la HBM.
Architecture unifiée : Mamba remplace l'alternance de blocs d'attention multi-têtes (MHA) et MLP par un bloc homogène unique qui combine la conception à portes de H3 avec le MLP à portes de type SwiGLU.

Méthode

Les modèles d'espace d'états continus projettent une entrée $x (t) \in R$ vers $y (t) \in R$ via un état latent implicite $h (t) \in R^{N}$ : $h^{'} (t) = A h (t) + B x (t), y (t) = C h (t)$

En utilisant une discrétisation par bloqueur d'ordre zéro (ZOH) avec un pas $Δ$ , le système continu se transforme en paramètres de récurrence discrets : $\overline{A} = exp (Δ A), \overline{B} = (Δ A)^{- 1} (exp (Δ A) - I) \cdot Δ B$ $h_{t} = \overline{A} h_{t - 1} + \overline{B} x_{t}, y_{t} = C h_{t}$

Le mécanisme de sélection (S6)

Au lieu de maintenir $(Δ, B, C)$ constants au cours du temps, Mamba les rend dépendants de l'entrée :

$B_{t} = s_{B} (x_{t}) = Linear_{N} (x_{t})$
$C_{t} = s_{C} (x_{t}) = Linear_{N} (x_{t})$
$Δ_{t} = τ_{Δ} (Parameter + s_{Δ} (x_{t}))$ , avec $τ_{Δ} = softplus$ et $s_{Δ} (x) = Broadcast_{D} (Linear_{1} (x))$

Cela généralise le mécanisme de porte des RNN classiques : lorsque $N = 1, A = - 1, B = 1$ , la récurrence se simplifie directement en une mise à jour à porte $h_{t} = (1 - g_{t}) h_{t - 1} + g_{t} x_{t}$ , où $g_{t} = σ (Linear (x_{t}))$ .

Expansion d'état adaptée au matériel

Pour préserver un débit élevé sans matérialiser l'état $(B, L, D, N)$ dans la HBM :

Les paramètres $(Δ, A, B, C)$ sont chargés depuis la HBM directement dans la SRAM intégrée.
La discrétisation en $(\overline{A}, \overline{B})$ et un scan associatif parallèle efficace en calcul sont exécutés entièrement dans la SRAM.
L'état latent est contracté avec $C$ vers la dimension $(B, L, D)$ avant d'être réécrit dans la HBM.
Les états intermédiaires sont recalculés dans la SRAM pendant la passe arrière au lieu d'être sauvegardés dans la HBM.

Résultats

Extrapolation synthétique : Sur la tâche des Induction Heads, Mamba entraîné sur une longueur de séquence de 256 extrapole parfaitement (100 % de précision) jusqu'à une longueur de $2^{20} = 1, 048, 576$ (soit $4000 \times$ la longueur de séquence), tandis que les modèles de référence à attention échouent ou provoquent un dépassement de mémoire (OOM) au-delà de $2^{14}$ . Sur le Selective Copying, Mamba atteint 99,8 % de précision contre 57,0 % pour H3 et 18,3 % pour S4.
Modélisation du langage : Évalué selon les lois d'échelle de Chinchilla jusqu'à 1,3B de paramètres sur the Pile, Mamba égale ou surpasse Transformer++ (plongements rotatifs, SwiGLU, RMSNorm). Sur les tâches aval zero-shot (ex. LAMBADA, PIQA, HellaSwag, WinoGrande, ARC), Mamba-2.8B atteint une précision moyenne de 63,3 %, surpassant Pythia-2.8B (59,1 %) et Pythia-6.9B (61,7 %), et égalant ou dépassant des modèles de 6B à 7B paramètres.
Génomique et audio : Sur la classification d'ADN de grands singes avec une longueur de séquence de 1M ( $2^{20}$ ), Mamba atteint 71,67 % (1,4M de paramètres) et 81,31 % (7M de paramètres) de précision contre 54,87 % pour HyenaDNA. Sur la génération de parole SC09, Mamba réduit le score FID de 1,99 (SaShiMi) à 0,94 (6,1M de paramètres) et 0,67 (24,3M de paramètres).
Efficacité : Le noyau de scan fusionné s'exécute 20 à 40× plus vite que les implémentations de scan standard de PyTorch et tourne jusqu'à 7× plus vite que FlashAttention-2 sur un contexte de 32K. Le débit d'inférence est jusqu'à 5× supérieur à celui de Transformers comparables.

Limites et questions ouvertes

Performances sur les signaux continus : Bien que la sélection dépendante de l'entrée améliore le raisonnement discret (texte, ADN), les ablations sur le préentraînement de formes d'onde audio brutes montrent que les modèles purement linéaires et invariants dans le temps (LTI) (comme le S4 standard) restent plus performants sur des données continues échantillonnées uniformément.
Échelle du modèle : Les modèles évalués plafonnent à environ 3B de paramètres et 300B de jetons de préentraînement ; les performances empiriques à des échelles de 7B–70B+ et au-delà nécessitent des recherches approfondies.
Écosystème aval : Les Transformers disposent d'un riche écosystème d'outils pour l'ingénierie d'invites (prompt engineering), la manipulation en contexte, la quantification du cache KV et le réglage par adaptateurs. L'adaptation des SSM à ces différents modes d'interaction demeure un champ d'exploration ouvert.

Termes clés

Structured State Space Models (SSMs): Classe de modèles de séquences profonds issus des systèmes dynamiques linéaires continus, associant des entrées 1D à travers un état latent implicite.
Linear Time Invariance (LTI): Propriété mathématique selon laquelle les paramètres de transition d'un système restent statiques à tous les pas de temps, permettant des représentations rapides par convolution globale mais restreignant le raisonnement dépendant du contenu.
Selection Mechanism: Paramétrage des dynamiques de transition du SSM (Δ, B, C) en fonction du jeton d'entrée, permettant au modèle de propager, réinitialiser ou ignorer dynamiquement le contexte le long de la séquence.
S6: Abréviation désignant les modèles d'espace d'états structurés sélectifs calculés via un scan associatif parallèle plutôt que par convolution.
Parallel Associative Scan: Algorithme de calcul de préfixes efficace en travail qui évalue une récurrence linéaire en parallèle sur la longueur de la séquence sur des accélérateurs matériels.
Hardware-Aware State Expansion: Stratégie de fusion de noyaux GPU et de recalcul effectuant la discrétisation et le scan dans la SRAM rapide sur puce afin d'éviter les allers-retours de grands états intermédiaires vers la mémoire à large bande passante (HBM).
Induction Heads: Mécanisme synthétique de rappel associatif servant à évaluer si un modèle est capable de reproduire un jeton précédemment observé après son jeton de contexte antérieur.
Mamba Block: Bloc de modèle de séquence homogène et simplifié intégrant projections d'entrée, convolution 1D, SSM sélectifs et portes multiplicatives, sans couches séparées d'attention ou de MLP.

Questions

Pourquoi les SSM sélectifs ne peuvent-ils pas être entraînés à l'aide de convolutions FFT rapides ?

Les convolutions globales rapides basées sur la FFT nécessitent que la dynamique du système (les matrices A, B et C) reste constante au cours du temps (invariance temporelle linéaire). Dans la mesure où Mamba fait de B, C et Δ des fonctions du jeton d'entrée à chaque pas de temps, le noyau de convolution varie à chaque position, ce qui brise l'équivalence convolutionnelle et impose l'utilisation de scans parallèles récurrents.

Comment Mamba évite-t-il l'explosion mémoire liée à l'expansion de l'état récurrent pendant l'entraînement ?

Une récurrence naïve avec une dimension d'expansion d'état N requiert la matérialisation d'un état intermédiaire de dimensions (B, L, D, N) dans la HBM du GPU. Mamba utilise la fusion de noyaux pour charger uniquement des entrées de faible dimension dans la SRAM, calcule la discrétisation et le scan parallèle au sein de la SRAM rapide, n'écrit que la sortie contractée (B, L, D) dans la HBM, et recalcule les activations intermédiaires lors de la rétropropagation.

Quelle est la relation entre le paramètre Δ de Mamba et le mécanisme de porte récurrent ?

Le paramètre de pas temporel Δ agit comme un mécanisme de porte généralisé. Discrétisée par bloqueur d'ordre zéro (ZOH) avec N=1 et A=-1, la récurrence sélective se réduit mathématiquement à une porte de RNN standard : un grand Δ incite le modèle à se concentrer sur l'entrée courante et à réinitialiser l'état, tandis qu'un petit Δ conserve l'état historique et ignore l'entrée.

Mamba utilise-t-il un cache KV pendant la génération ?

Non. Mamba fonctionnant comme une récurrence en espace d'états autorégressive lors de la génération, il ne conserve qu'un état caché de taille fixe (de dimension D × N par couche) au lieu de mettre en cache les clés et valeurs de tous les jetons passés, ce qui permet une inférence par pas en temps constant O(1).

La présentation de cette page a été rédigée par une IA à partir de l'article et peut contenir des erreurs. Consultez l'article pour les détails.

Lire avec traduction et IA