Mamba et les architectures State Space Models : l'alternative aux Transformers ?
Depuis 2017, les Transformers dominent le paysage de l'intelligence artificielle, des grands modèles de langage (LLM) à la vision par ordinateur. Pourtant, une nouvelle famille d'architectures commence à sérieusement grignoter leur monopole : les State Space Models (SSM), et particulièrement le modèle Mamba. Promettant une complexité linéaire contre une complexité quadratique pour les Transformers, ces modèles pourraient bien résoudre le goulot d'étranglement de la fenêtre de contexte. Dans cet article, nous allons décortiquer le fonctionnement des SSM, comprendre les innovations de Mamba, analyser les benchmarks et voir comment implémenter ces modèles en pratique.
Prérequis
- Notions fondamentales en Deep Learning (réseaux de neurones, rétropropagation)
- Compréhension de base de l'architecture Transformer (mécanisme d'attention, self-attention)
- Connaissances en algèbre linéaire (multiplication de matrices, vecteurs d'état)
- Bases de Python et de PyTorch
- Notions sur la complexité algorithmique (notation Grand O)
Le problème des Transformers : le coût de l'attention
Pour comprendre l'essor des SSM, il faut d'abord comprendre les limites actuelles des Transformers. Le cœur d'un Transformer est le mécanisme d'attention, qui permet au modèle de mettre en relation chaque token d'une séquence avec tous les autres tokens. Mathématiquement, pour une séquence de longueur $L$ et une dimension de modèle $d$, la matrice d'attention a une taille de $L \times L$.
Cela conduit à une complexité temporelle et spatiale de $O(L^2 \cdot d)$. En d'autres termes, si vous doublez la taille de votre contexte (par exemple, passer de 32k à 64k tokens), le temps de calcul et la mémoire requise sont multipliés par quatre. C'est ce qu'on appelle la complexité quadratique.
Pour visualiser cet écart, un outil de tracé comme Matplotlib permet de comparer la courbe de croissance quadratique des Transformers avec la courbe de croissance linéaire des State Space Models. En échelle logarithmique, la courbe du Transformer s'envole très rapidement à mesure que la longueur de la séquence augmente, tandis que celle du SSM reste remarquablement plate, illustrant l'avantage majeur de ces nouvelles architectures sur les longues séquences.
Le bottleneck de la fenêtre de contexte
Cette limitation quadratique force les créateurs de LLMs à tronquer les longs textes. Analyser un livre entier, un dépôt GitHub complet ou de longues vidéos nécessite des techniques de contournement lourdes (comme la compression de contexte ou les architectures hiérarchiques) qui dégradent la qualité des réponses. Nous avons besoin d'un mécanisme capable de traiter des séquences infinies avec un coût fixe par token. C'est exactement ce que proposent les State Space Models.
Introduction aux State Space Models (SSM)
Les State Space Models puisent leurs racines dans la théorie du contrôle, un domaine de l'ingénierie et des mathématiques utilisé pour modéliser des systèmes physiques continus (comme le vol d'un avion ou le circuit électrique d'un robot). L'idée est de modéliser un système en observant comment un signal d'entrée $u(t)$ modifie un état interne caché $x(t)$ pour produire un signal de sortie $y(t)$.
Les équations d'état continues
Un SSM continu est défini par deux équations différentielles linéaires :
$$x'(t) = Ax(t) + Bu(t)$$
$$y(t) = Cx(t) + Du(t)$$
Où :
- $x(t)$ est le vecteur d'état latent (la "mémoire" du système)
- $u(t)$ est le signal d'entrée (par exemple, un flux de texte continu)
- $y(t)$ est le signal de sortie
- $A$, $B$, $C$, $D$ sont des matrices de paramètres apprenables
- $A$ est particulièrement importante : c'est la matrice de transition d'état (State Transition Matrix). Elle détermine comment l'état évolue naturellement au fil du temps.
La discrétisation
Le problème, c'est que nous travaillons en informatique avec des données discrètes (des tokens de texte, des pixels). Nous devons convertir ces équations continues en équations discrètes par pas de temps $\Delta$. L'article fondateur S4 (Structured State Spaces for Sequence Modeling) a introduit une méthode de discrétisation élégante basée sur la transformation bilinéaire de l'approximation de Zéro-Order Hold (ZOH) :
$$\bar{A} = (I - \Delta A)^{-1}$$
$$\bar{B} = (I - \Delta A)^{-1} \Delta B$$
Cela nous donne la récurrence discrète :
$$x_k = \bar{A}x_{k-1} + \bar{B}u_k$$
$$y_k = Cx_k$$
Pourquoi c'est puissant ?
Cette équation récursive $x_k = \bar{A}x_{k-1} + \bar{B}u_k$ a une propriété magique en informatique : elle s'effectue en temps constant $O(1)$ par token. Peu importe la longueur de la séquence, pour calculer l'état actuel, vous n'avez besoin que de l'état précédent et du token actuel. La complexité globale passe donc à $O(L)$ ! De plus, grâce à la propriété de convolution, ces opérations peuvent être parallélisées sur GPU lors de l'entraînement, tout en offrant un décodage récursif ultra-rapide lors de l'inférence.
Le problème du contenu sélectif : l'arrivée de Mamba
Si les SSM sont si puissants, pourquoi n'ont-ils pas remplacé les Transformers plus tôt ? Le problème majeur des SSM classiques (comme S4) réside dans leur inability à faire du contenu sélectif (selective copy/paste).
Dans un texte, certains mots sont du remplissage ("le", "et", "un"), tandis que d'autres portent une information critique ("clé API", "mot de passe", un nom propre). Le mécanisme d'attention des Transformers excelle à ça : il peut décider d'envoyer 100% de son attention sur un token spécifique, même s'il est situé 10 000 tokens plus loin.
Les SSM classiques, en revanche, traitent chaque token avec les mêmes paramètres $\bar{A}$, $\bar{B}$, $C$. La matrice $\bar{A}$ est généralement initialisée pour être une matrice de type HiPPO (High-order Polynomial Projection Operators), conçue pour retenir l'historique global de la séquence comme un polynôme. C'est excellent pour résumer, mais terrible pour ignorer le bruit et se concentrer soudainement sur un détail précis.
L'innovation de Mamba : des paramètres dépendants de l'entrée
C'est ici qu'intervient Albert Gu et Tri Dao (créateur de FlashAttention) avec leur papier fondamental Mamba: Linear-Time Sequence Modeling with Selective State Spaces (arxiv.org/abs/2312.00752). Leur idée est de rendre les paramètres du SSM dynamiques : ils dépendent de l'entrée $u_k$.
- $B_k, C_k = \text{Linear}_{BC}(u_k)$
- $\Delta_k = \text{Softplus}(\text{Linear}_{\Delta}(u_k))$ (Le pas de temps $\Delta$ devient également dynamique !)
Si le modèle reconnaît un token important, il peut modifier $\Delta_k$ pour "ralentir" le temps, permettant au token de s'imprimer profondément dans la matrice d'état $x_k$. S'il s'agit de bruit, il "accélère" le temps, et le token est oublié presque instantanément.
En pratique avec PyTorch, cette logique se traduit par l'ajout de modules nn.Linear qui projettent le token d'entrée vers les vecteurs $B$, $C$ et $\Delta$. La fonction d'activation Softplus est appliquée sur $\Delta$ pour garantir que le pas de temps reste toujours positif, stabilisant ainsi la mise à jour de l'état caché au fil de la séquence.
Le verrou matériel (Hardware-aware algorithm)
Rendre $B$, $C$ et $\Delta$ dynamiques casse la propriété de convolution. On ne peut plus utiliser la FFT (Fast Fourier Transform) pour paralléliser le calcul sur GPU. On est forcé de revenir à la boucle récursive séquentielle, ce qui est extrêmement lent sur les GPU modernes conçus pour le calcul massivement parallèle.
C'est la deuxième géniale innovation de Mamba. Les auteurs ont conçu un algorithme spécifique au niveau du matériel (inspiré de FlashAttention) :
1. Au lieu de stocker l'état en mémoire globale (VRAM), ils le stockent dans la mémoire ultra-rapide ultra-localisée du GPU (SRAM).
2. Ils utilisent un mécanisme de "scan" qui effectue la récurrence de manière parallèle par blocs ( associative scan ), exploitant les warp-level primitives des GPU.
Cela permet à Mamba de conserver la complexité temporelle $O(L)$ tout en atteignant des vitesses d'entraînement comparables, voire supérieures, aux Transformers.
Architecture détaillée de Mamba
Un bloc Mamba complet fonctionne comme un empilement intelligent de différentes couches conçu pour traiter la séquence de manière hiérarchique. Voici comment se déroule le passage d'un token à travers un bloc Mamba :
1. Expansion de la dimension
La séquence d'entrée de dimension $d_{model}$ (par exemple 768) est d'abord projetée dans un espace de dimension plus élevé $E \times d_{model}$ (souvent $E=2$) via une convolution de noyau de taille $K$ (généralement $K=4$). Cette convolution agit comme un pré-filtre local, permettant au modèle de regarder les $K$ tokens précédents immédiats, ce que la récurrence globale ne fait pas nativement.
2. Sélection dynamique
C'est ici qu'interviennent les projections linéaires que nous avons vues plus haut pour générer $B_k$, $C_k$ et le pas de temps $\Delta_k$. C'est le cœur du "Selective State Space".
3. Le Scan SSM
L'algorithme de scan matériel applique la récurrence discrète en utilisant les paramètres dynamiques pour mettre à jour l'état caché $x_k$.
4. Projection de sortie et connexion résiduelle
La sortie du scan est projetée de l'espace étendu vers l'espace $d_{model}$ d'origine. Une connexion résiduelle (skip connection) est ajoutée, similaire à ce qui se fait dans les Transformers, combinée à une normalisation.
L'implémentation d'un bloc Mamba complet en PyTorch s'articule autour de plusieurs composants clés : une projection d'entrée suivie d'une convolution causale 1D pour capturer le contexte local, puis les calculs des paramètres dynamiques via des couches linéaires. L'état caché est initialisé selon la structure HiPPO, et la récurrence SSM est finalement résolue via une opération de scan optimisée (généralement un kernel CUDA sur mesure pour des raisons de performance). La sortie du scan est ensuite combinée à une branche de gating (multiplication par une porte d'activation) avant la projection finale.
Mamba-2 : Structured State Space Duality
Si Mamba a marqué une étape en 2023, la sortie de Mamba-2 en 2024 a consolidé cette architecture en lui donnant des fondations mathématiques encore plus solides. Tri Dao et Albert Gu ont prouvé dans leur nouveau papier que les SSM et l'Attention partagent en fait une structure duale commune basée sur l'algèbre linéaire structurée.
Le passage à l'espace des têtes (Head-wise)
Mamba-2 modifie la façon dont la matrice d'état $A$ et les paramètres $B, C$ sont structurés. Au lieu de traiter l'espace latent comme un grand bloc monolithique, Mamba-2 le divise en "têtes" (heads), exactement comme les Multi-Head Attention des Transformers.
La récurrence devient :
$$x_k^h = A^h x_{k-1}^h + B_k^h u_k^h$$
L'algorithme SDL (Structured Dot Product)
La grande innovation de Mamba-2 est le remplacement du associative scan par un algorithme appelé Structured Dot Product (SDL). Au lieu de faire une multiplication par une matrice d'état pleine, Mamba-2 exploite la structure de la matrice $A$ (souvent structurée comme une matrice de Toeplitz ou de Hankel normalisée) pour transformer la récurrence en un produit scalaire structuré.
Cela permet à Mamba-2 d'utiliser les mêmes noyaux d'optimisation mémoire (tiling, flash-like memory IO) que FlashAttention-2, résultant en une accélération de l'entraînement de 2 à 8 fois par rapport à Mamba-1, tout en permettant de gérer des dimensions d'état beaucoup plus grandes.
Résultats et Benchmarks
La théorie est séduisante, mais qu'en est-il des performances réelles ? C'est ici que Mamba justifie tout le battage médiatique.
Efficacité matérielle (Throughput)
C'est le domaine où Mamba écrase littéralement les Transformers. Grâce à sa complexité $O(L)$ et à son algorithme optimisé pour le matériel, Mamba ne souffre pas de l'explosion mémoire lors de l'augmentation de la fenêtre de contexte.
- En inférence (génération de tokens) : Mamba est entre 4 et 5 fois plus rapide qu'un Transformer équivalent pour des séquences longues, car il n'a pas besoin de stocker et de recalculer le cache KV (Key-Value cache) pour tous les tokens précédents.
- En entraînement : Mamba atteint des débits (tokens par seconde par GPU) nettement supérieurs lorsque la longueur de séquence dépasse 8192 tokens.
Performances sur les tâches de modélisation de langage
Sur les benchmarks classiques de NLP (Pile, WikiText, LAMBADA), les modèles de type Mamba (comme la famille de modèles ouverts fournie par les auteurs) rivalisent directement avec les Transformers de même taille.
- Modèles de petite à moyenne taille (< 1B paramètres) : Mamba surpasse ou égale les Transformers (comme S4 ou les petits Llama) en perplexité.
- Modèles à grande échelle (> 2.5B paramètres) : C'est le seul point où les Transformers (type Llama ou Mistral) gardent un léger avantage sur les tâches de raisonnement pur. L'attention globale reste légèrement plus puissante pour la résolution de problèmes logiques complexes nécessitant de faire des sauts multiples dans le texte.
Le triomphe sur le "Selective Copying"
Mamba a été spécifiquement conçu pour résoudre les tâches que les anciens SSM échouaient à résoudre. Sur les tâches synthétiques de copie sélective (par exemple, "Trouve la clé API cachée dans ce texte de 50 000 mots et recopie-la"), Mamba atteint un taux de succès de 100%, là où les SSM classiques tombent à 0% et où les Transformers souffrent d'un fort ralentissement.
Hybrides : le meilleur des deux mondes ?
La tendance actuelle dans l'industrie (vue chez Microsoft, Mistral, ou dans des architectures comme Jamba) n'est pas de tuer les Transformers, mais de les hybrider.
L'architecture Jamba (d'AI21 Labs) empile des couches Mamba avec des couches d'Attention. Par exemple, un modèle peut utiliser 8 couches Mamba pour 1 couche d'Attention. Le Mamba absorbe 95% du contexte long et du bruit de manière ultra-rapide, et l'Attention est utilisée ponctuellement pour faire les connexions complexes et le raisonnement profond. Ce type d'hybride offre un ratio qualité/vitesse exceptionnel et une fenêtre de contexte atteignant 256000 tokens avec une empreinte VRAM divisée par 4. D'ailleurs, si vous hésitez sur l'approche à privilégier pour vos projets, notre guide sur le RAG vs fine-tuning vs agents : choisir la bonne approche en 2026 peut vous aider à y voir plus clair.
Comment utiliser Mamba en pratique
L'écosystème autour de Mamba grandit rapidement. La façon la plus simple de l'expérimenter est d'utiliser la bibliothèque officielle mamba-ssm, développée en grande partie par Tri Dao.
Installation
L'installation requiert une compilation CUDA locale, car le cœur de Mamba est écrit en C++/Triton pour des raisons de performance. Il est recommandé de démarrer dans un environnement virtuel isolé en utilisant conda. Une fois l'environnement activé avec Python 3.10, la procédure consiste à installer PyTorch en spécifiant la version 12.1 de CUDA, puis à ajouter les deux dépendances obligatoires que sont causal-conv1d et mamba-ssm via pip.
Utilisation basique avec Hugging Face
De nombreux modèles pré-entraînés basés sur Mamba sont maintenant disponibles sur le Hub Hugging Face. Vous pouvez les utiliser exactement comme un modèle Transformer classique grâce à la librairie transformers. Le chargement s'effectue avec AutoModelForCausalLM et AutoTokenizer en spécifiant l'identifiant du modèle (comme state-spaces/mamba-2.8b). Après avoir tokenisé votre prompt, il suffit d'appeler la méthode generate en précisant des paramètres classiques comme max_new_tokens (fixé autour de 200 pour une réponse courte), une temperature (autour de 0.7) et l'activation de l'échantillonnage, avant de décoder le résultat obtenu.
Implémentation d'une couche Mamba personnalisée
Si vous construisez votre propre architecture from scratch (par exemple pour un modèle de classification de séries temporelles ou de génomique, domaines où Mamba excelle), vous interagirez directement avec l'opération de base. La bibliothèque mamba-ssm expose un module Mamba prêt à l'emploi. Pour créer un classifieur de séquences, il suffit d'empiler plusieurs de ces modules dans une nn.ModuleList, en leur passant les hyperparamètres clés comme la dimension du modèle (par exemple 256 ou 512), la dimension de l'état SSM (16), la taille du noyau de convolution local (4) et le facteur d'expansion (2). Chaque bloc est encapsulé dans une connexion résiduelle, suivie d'une normalisation RMSNorm. En sortie, une simple couche linéaire appliquée sur le dernier token de la séquence permet d'obtenir les logits de classification.
Récapitulatif
- Complexité algorithmique : Les SSM comme Mamba fonctionnent en $O(L)$ (linéaire) comparé au $O(L^2)$ (quadratique) des Transformers, permettant de traiter des séquences quasi-infinies sans saturation de la VRAM.
- Mécanisme sélectif : Mamba résout le défaut historique des SSM en rendant ses paramètres ($B$, $C$, $\Delta$) dynamiques et dépendants du token d'entrée, enabling un filtrage précis de l'information.
- Optimisation matérielle : Grâce à des algorithmes inspirés de FlashAttention (associative scan dans Mamba-1, SDL dans Mamba-2), la théorie mathématique se traduit par de réelles accélérations sur GPU.
- Efficacité en inférence : Mamba élimine le besoin d'un cache KV (Key-Value) coûteux en mémoire, rendant la génération de texte beaucoup plus rapide et économe.
- Performances : Sur les benchmarks de langage et de copie sélective, Mamba rivalise ou surpasse les Transformers à taille égale, bien que les très grands Transformers gardent un léger avantage sur le raisonnement pur.
- Tendance Hybride : L'avenir probable n'est pas un remplacement total, mais une hybridation (ex: Jamba) combinant des couches Mamba pour le contexte long et ponctuellement des couches d'Attention pour le raisonnement.
Conclusion
Mamba et les architectures State Space Models représentent la première véritable alternative paradigmatique aux Transformers depuis 2017. En s'attaquant au problème fondamental de la complexité quadratique tout en résolvant le problème du contenu sélectif, Mamba redéfinit les limites de ce que nous pouvons faire avec le traitement de séquences. Si les très grands LLMs actuels continuent de dominer les benchmarks de raisonnement général en 2025, l'efficacité brute, la réduction des coûts d'inférence et la capacité à ingérer des contextes gigantesques font des SSM l'outil incontournable des prochaines années, en particulier pour les architectures multimodales (où la vidéo et l'audio nécessitent des fenêtres de contexte colossales).
Prêt à passer à l'action et à intégrer ces nouvelles architectures dans vos workflows ? Rejoignez la communauté d'AI-Master.dev pour accéder à nos ateliers pratiques, télécharger des notebooks optimisés et maîtriser les modèles de demain.
[S'abonner à la newsletter AI-Master.dev]