Qu’est-ce qu’un modèle Mamba ?

Les modèles Mamba sont probablement les premiers modèles de deep learning capables de rivaliser avec les transformers sur la tâche qui a fait leur renommée : la modélisation du langage. Plus précisément, l’architecture Mamba a démontré une capacité à égaler les transformers de taille équivalente sur les principaux benchmarks LLM, tout en étant nettement plus efficace en termes de latence et d’utilisation de la mémoire.

L’architecture Mamba a été introduite pour la première fois par Tri Dao et Albert Gu dans l’article de 2023, « Mamba: Linear-Time Sequence Modeling with Selective State Spaces ». Un an plus tard, ils ont publié un second article explorant davantage les liens entre SSM et transformers et présentant une version affinée et nettement plus rapide de cette architecture : Mamba-2.

Même si les transformers sont restés le paradigme dominant dans les LLM pendant les deux années ayant suivi la sortie de l’article Mamba original, l’architecture a été intégrée à un nombre croissant de modèles open source. Certains, comme le Codestral Mamba de Mistral AI, sont des modèles 100 % Mamba. D’autres, comme la série Jamba (AI21) ou Granite 4.0 (IBM), sont des modèles hybrides combinant des couches d’attention (transformer) et des couches SSM (Mamba). En plus de leurs avantages en termes de performances, la prolifération des modèles Mamba contribue à démocratiser l’accès à l’IA, puisqu’ils peuvent tourner de manière fluide sur du matériel beaucoup moins coûteux.

À l’origine, les SSM ont été conçus pour prédire l’état suivant d’une séquence continue (signal électrique, conditions météorologiques ou trajectoire d’un objet en mouvement) à partir d’une entrée. Sur le plan conceptuel comme sur le plan mathématique, ils sont proches des réseaux neuronaux récurrents (RNN) qui dominaient le traitement automatique du langage naturel (NLP) avant l’arrivée des transformers en 2017, et d’autres algorithmes de machine learning comme les réseaux neuronaux convolutifs (CNN) et les modèles de Markov cachés (HMM).

Comme leur nom l’indique, les SSM font des prédictions sur l’état suivant dans un système dynamique en modélisant l’espace d’état : une représentation mathématique de toutes les variables d’état qui décrivent l’état d’un système, et la plage de valeurs possibles pour chacune de ces variables, associées les unes avec les autres.

Un SSM prend une séquence d’entrée x(t) et la mappe à un état latent h(t), l’équivalent de l’état caché d’un RNN, afin de prédire une séquence de sortie y(t). Au cœur de tout SSM, on trouve deux équations :

• L’équation d’état,  $h^{\prime }\left(t\right)=A*h\left(t\right)+B*x\left(t\right)$ 
• L’équation de sortie,  $y\left(t\right)=C*h\left(t\right)+D*x\left(t\right)$ 

Les paramètres clés du modèle sont les matrices A, B, C et D, généralement représentées sous forme de matrices de poids. Dans les domaines où les SSM sont traditionnellement utilisés, comme la théorie du contrôle, ces matrices sont généralement supposées fixes : elles représentent les dynamiques d’un système établi, et le SSM est utilisé pour trouver les entrées x qui conduisent à des sorties y souhaitables.Dans les SSM modernes, ces matrices sont elles-mêmes des paramètres à optimiser via machine learning. Dans les modèles de deep learning, ces matrices sont représentées par les poids apprenables d’un réseau de neurones.

L’équation d’état décrit l’évolution de l’état. Les valeurs de la matrice A déterminent comment chaque variable d’état évolue dans le temps sans influence externe. Les valeurs de la matrice B déterminent comment l’entrée, comme le token suivant dans une séquence de texte, influence chaque variable d’état.

En modélisation du langage, l’état courant représente le contexte d’une séquence de texte, mis à jour après chaque token. Son rôle est équivalent à celui du cache KV dans un transformer.

L’équation de sortie décrit comment l’état courant influence la sortie (via la matrice C) ainsi que la manière dont l’entrée influence directement la sortie (via la matrice D). Étant donné que la matrice D ne fait pas réellement partie de la modélisation de h(t) elle-même, elle est souvent omise dans les diagrammes et les discussions sur les SSM, pour se concentrer sur les matrices principales A, B et C.

Dans un LLM Mamba, l’équation de sortie est utilisée pour générer le token suivant.

Les SSM traditionnels sont conçus pour modéliser des entrées continues, mais les séquences de texte (et la plupart des autres modalités de données traitées par les modèles de deep learning modernes) sont des entrées discrètes. Utiliser un SSM pour modéliser une séquence discrète suppose qu’on représente ses pas de temps distincts et spécifiques comme faisant partie d’un signal continu.

Conceptuellement, la discrétisation revient à échantillonner la valeur d’une fonction continue à des moments précis. Cela introduit un nouveau paramètre, la taille de pas, représentée par ∆, qui détermine combien de temps on échantillonne ou « maintient » cet valeur à chaque pas de temps discret t. Changer ∆ revient un peu à modifier la résolution des données (de séries temporelles), ou la fréquence d’images (d’une vidéo). Il existe plusieurs méthodes de « discrétisation », mais la plupart des variantes de SSM modernes (y compris Mamba) utilisent la méthode simple du maintien d’ordre zéro (ZOH).

La discrétisation d’un SSM lui permet d’être utilisé comme un RNN pour des tâches séquence à séquence. Les équations et paramètres d’un SSM discrétisé sont généralement réécrits en utilisant la notation à indices propre aux RNN pour les distinguer de leurs équivalents en temps continu. Dans cette notation, h_t représente l’espace d’état mis à jour que le modèle va générer et h_t-1 représente l’état qui le précède, c’est-à-dire l’espace d’état actuel.

 $h_t=\overline{A}{h}_{t-1}+\overline{B}{x}_t$ 
 $y_t=\overline{C}{h}_t$

La modélisation de données textuelles à l’aide de SSM discrets standard est peu commode, car ils héritent des mêmes limites que les RNN. Deux de ces inconvénients ont été corrigés par l’introduction de modèles de séquence d’espace d’état structuré (ou « modèles S4 » pour Structured State Space Sequence) par Albert Gu et al. en 2021 : leur inefficacité inhérente lors de l’entraînement des modèles et leur incapacité à gérer les longues séquences.

Cependant, le succès des modèles S4, et de leurs nombreux dérivés, tels que les DSS (Diagonal SSM), S4D (Diagonal S4) et H3, a directement ouvert la voie à ce qui est devenu Mamba.

L’avantage des SSM discrétisés, qui sont alors l’équivalent d’une instance spécifique d’un RNN, c’est que les RNN sont extrêmement rapides en inférence. L’inconvénient, en revanche, c’est que les RNN sont très lents à entraîner.

Heureusement, les SSM discrétisés possèdent une propriété qui les distingue des autres RNN : ils ne modélisent que des dépendances linéaires. En d’autres termes, il utilisent uniquement des additions et des multiplications simples. Comme le montre l’article S4, ces récurrences linéaires (simples, répétées et interdépendantes) peuvent être déroulées sous la forme d’un noyau de convolution à une dimension, qui mappe directement l’entrée x à la sortie y en une seule étape : Ce calcul peut être réalisé très efficacement grâce à la transformée de Fourier rapide (FFT).

La seule « limite », c’est que cela n’est possible que lorsque l’intégralité de la séquence d’entrée est connue. Ce n’est pas possible pendant l’inférence, mais c’est possible pendant l’entraînement. Cela permet à un SSM structuré de profiter du meilleur des deux mondes : en entraînement, il fonctionne très efficacement comme un CNN ; en inférence, il fonctionne très efficacement comme un RNN.

Comme la plupart des RNN, les SSM standard sont intrinsèquement peu performants pour modéliser les dépendances à longue distance. En d’autres termes, ils ne sont pas capables de comprendre la relation entre les étapes d’une séquence très éloignées les unes des autres, comme les mots au début et à la fin d’un paragraphe, ce qui les rend incapables de modéliser de longues séquences.

Pour résoudre ce problème, Gu et ses co-auteurs (dont Tri Dao) ont utilisé une technique appelée HiPPO (High-order Polynomial Projection Operators), qui consiste à structurer les valeurs initiales des matrices A et B à partir de formules issues des polynômes orthogonaux. Cela contraste avec la pratique standard de machine learning, dans laquelle les poids des modèles sont initialisés de manière aléatoire au début de leur entraînement. Dans S4, Tri Dao et Albert Gu utilisent des schémas d’initialisation dérivés des polynômes de Legendre, et explorent d’autres formules dans un autre article intitulé « How to Train Your HiPPO ».¹

L’article S4 indique que « le simple fait de remplacer une matrice A aléatoire par [la matrice HiPPO] a permis de faire passer la performance sur le benchmark MNIST séquentiel de 60 % à 98 % », résolvant ainsi le problème de mémoire à long terme des SSM. Les variantes ultérieures des SSM structurés (DSS, S5 ou Mamba) utilisent des schémas d’initialisation différents (souvent plus simples) pour A et B, mais conservent les principes fondamentaux de la méthode HiPPO : une structure diagonale qui impose des mises à jour stables et une certaine indépendance entre les valeurs de la matrice.

Au cœur de l’architecture Mamba, on trouve deux innovations clés : La première est le modèle d’espace d’état sélectif, qui confère à Mamba une capacité essentielle, jusque-là propre aux transformers : la capacité à se concentrer sur certaines parties de l’historique des entrées passées et à en ignorer d’autres, en fonction de leur pertinence actuelle. La seconde est le parallel scan optimisé matériellement, un algorithme qui organise les calculs dans la hiérarchie mémoire du processeur graphique (GPU) de manière à maximiser leur vitesse et leur efficacité.

Dans un transformer, cette capacité est assurée par le mécanisme d’attention, qui ajuste les poids d’attention pour renforcer (ou diminuer) l’influence de chaque token précédent en fonction de sa pertinence pour le token d’entrée courant. Les SSM classiques sont conçus pour utiliser tout l’historique des entrées afin de produire une sortie.Cela peut être acceptable (voire souhaitable) pour certaines tâches de modélisation, mais constitue un handicap majeur pour les tâches avancées de modélisation du langage.

Pour pallier cette incapacité à ignorer ou mettre en avant certaines parties de l’historique des entrées, Tri Dao et Albert Gu ont introduit une nouvelle classe de SSM dotée d’un sélecteur, le « selective scan ». Dans l’article Mamba, les auteurs précisent qu’ils « abrègent parfois les SSM sélectifs sous le terme modèles S6, car ce sont des modèles S4 dotés d’un mécanisme de sélection et calculés par scan ». Ils ont surnommé leur architecture S6 « Mamba » parce que tous ces S rappellent le sifflement d’un serpent.

Mamba peut donc être vu comme une architecture de réseau neuronal qui contient, en son cœur, un modèle d’espace d’état sélectif. Pour faire une analogie simple, Mamba est aux SSM sélectifs ce que le transformer est au mécanisme d’attention.

Dans un SSM traditionnel, les dynamiques sont fixes : les règles qui gouvernent l’évolution de l’état caché (les paramètres du modèle) sont les mêmes pour toutes les entrées et à chaque étape de la séquence. C’est ce qu’on appelle l’invariance dans le temps (LTI). Pour donner aux SSM la capacité de favoriser ou d’ignorer certaines parties du passé en fonction du contexte, Tri Dao et Albert Gu ont reconfiguré leur modèle de sorte que les valeurs des paramètres clés du modèle soient différentes pour différentes entrées.

Plus précisément, les SSM sélectifs font de la taille de pas ∆_t et des matrices B_t et C_t des fonctions directes du token d’entrée actuel x_t. Pour ce faire, on doit d’abord passer l’embedding vectoriel de _xt dans trois couches de projection linéaire parallèles, en d’autres termes, les couches de réseaux neuronaux à propagation avant standard (ou couches MLP). Cela équivaut à la manière dont les têtes query, key et value parallèles génèrent les vecteurs Q, K et V d’une entrée dans un modèle transformer.

En multipliant l’embedding du token x_t par les poids et biais de ce réseau de projection linéaire, on obtient donc les valeurs de ∆_t, B_t et C_t. Les poids et biais des couches de projection linéaire elles-mêmes sont appris pendant une phase de pré-entraînement sur de très grands jeux de données de corpus de texte, puis (éventuellement) ajustés par réglage fin.

• La valeur de ∆_t détermine l’importance de l’influence de x_t sur la mémoire du contexte accumulé jusque-là : en d’autres termes, elle contrôle l’ampleur de la mise à jour par rapport entre l’état caché h_t-1 et h_t. Une taille de pas ∆ plus élevée entraîne une plus grande modification de l’état et accélère la décroissance (« l’oubli ») des informations plus anciennes contenues dans l’état. Inversement, une taille de pas plus faible entraîne une mise à jour plus petite. Si le pas est suffisamment petit, l’entrée courante n’aura aucun impact sur l’état caché.
  
  
• Les modifications apportées à la matrice B_k déterminent comment le token d’entrée courant modifie l’état caché. Par exemple, si le token x_t représente le mot « hier », B_t pourrait être ajustée de façon à mettre l’état à jour pour refléter le fait que le contexte à venir concerne probablement le passé.
  
  
• Les modifications apportées à la matrice C_t déterminent comment ces informations contextuelles influencent la sortie du modèle y_t. En reprenant le même exemple, si le token x_k représente le mot « hier », C_t pourrait être modifiée de sorte que tous les verbes générés ensuite par le modèle soient conjugués au passé.

À noter : aucune adaptation dépendant de l’entrée n’est appliquée à la matrice A. Son rôle reste le même que dans les modèles S4 : mémoriser efficacement l’historique complet des entrées passées. Les matrices B et C déterminent quelles parties de cet historique utiliser à un moment donné.

Mais dès lors que le modèle n’est plus invariant dans le temps, il ne peut plus utiliser l’astuce de convolution pendant l’entraînement puisque le noyau de transition n’est plus constant : au cœur du mécanisme de sélectivité, la transition de h_t-1 vers h_t dépend désormais du contexte.

Pour contourner cette limite, Mamba utilise une astuce élégante permettant d’obtenir des bénéfices de parallélisation similaires. Comme les calculs du SSM ne reposent que sur des additions et des multiplications, ils respectent la propriété associative : on peut regrouper les opérations dans n’importe quel ordre sans modifier le résultat final. Cela permet de décomposer les nombreux calculs séquentiels en petits blocs indépendants qui peuvent être traités en parallèle par un GPU via un parallel prefix-sum scan.

De plus, les résultats sont combinés selon une hiérarchie spécifique exploitant de manière optimale les différents types de mémoire matérielle du GPU, selon des principes similaires aux techniques FlashAttention, également développées par Tri Dao, qui sont désormais omniprésentes dans les LLM modernes.  

Dans l’architecture Mamba, le modèle S6 fait office de module à l’intérieur d’un « bloc Mamba » plus large, de la même façon que le mécanisme d’attention fait office de module au sein d’un « bloc d’attention  plus large. Le module S6 est combiné à une architecture de réseau neuronal à portes. Les modèles Mamba comprennent généralement plusieurs blocs Mamba, c’est-à-dire une série de couches Mamba consécutives dans un réseau de neurones (avant la couche de sortie qui effectue la prédiction finale du modèle).

Avant d’entrer dans le bloc Mamba, une copie de l’entrée est envoyée directement à la sortie : c’est la connexion résiduelle. Le rôle du bloc Mamba est non seulement d’identifier quelles parties du contexte plus global sont pertinentes pour cette entrée, mais aussi de déterminer dans quelle mesure ce contexte devrait modifier le sens d’origine de l’entrée.

À l’intérieur du bloc Mamba, le vecteur d’entrée d’origine est traité comme suit :

• L’entrée est d’abord passée dans une couche linéaire deux fois plus large que le vecteur d’entrée d’origine afin de le projeter dans un espace de plus grande dimension. Par exemple, si le modèle représente à l’origine chaque token d’entrée x par un embedding de 512 dimensions, la multiplication de x par les poids de la couche de projection linéaire le transforme en embedding de 1 024 dimensions.
  
  
• Le vecteur élargi est scindé en deux. Une moitié (que nous appellerons x_proj) est envoyée dans le chemin qui passe par le SSM. L’autre moitié (que nous appellerons z_proj) est envoyée dans un chemin distinct qui passe par un mécanisme de gating. Pour plus de clarté dans les schémas, cette étape est souvent représentée comme deux couches linéaires parallèles.
  
  
• Avant d’atteindre le SSM, x_proj passe par une couche de convolution à une dimension. Cette couche de convolution extrait les patterns locaux (tels que les dépendances entre les tokens voisins, comme les paires verbe-sujet simples). Cela permet au SSM de « se concentrer » sur la compréhension contextuelle des dépendances globales à longue distance.
  
  
• La sortie de la couche de convolution sert d’entrée à une fonction d’activation non linéaire. L’introduction de la non-linéarité est l’une caractéristique fondamentale de tous les réseaux de neurones : elle leur permet de capturer des schémas plus complexes. Dans l’article Mamba, l’activation utilisée est le Sigmoid Linear Unit (SiLU). On appellera le vecteur résultant x_act.
  
  
• En parallèle, dans le chemin du mécanisme de gating, z_proj est également soumis à une fonction d’activation non linéaire, produisant z_act.
  
  
• Dans le chemin SSM, x_act est envoyé dans trois couches de projection linéaire parallèles qui génèrent respectivement les valeurs de ∆_x, B_x et C_x.
  
  
• Le SSM utilise ensuite ces paramètres dépendants de l’entrée (ainsi que les matrices A et D) pour calculer la mise à jour de l’espace d’état et produire la sortie du SSM : y.
  
  
• Le vecteur de sortie du SSM, y , est multiplié élément par élément par le vecteur de sortie z_act du chemin de gating. En pratique, chaque élément de z_act agit comme un potentiomètre sur une console de mixage audio : si un élément donné de z_act est proche de zéro, la multiplication par la partie correspondante de y donnera une valeur plus proche de zéro et son influence sera atténuée. Inversement, si un élément donné de zact est élevé, la multiplication par y amplifiera l’influence de ses informations contextuelles.
  
  
• Le vecteur résultant est ensuite projeté de nouveau vers sa taille d’origine. On peut l’interpréter comme un vecteur de mises à jour contextuelles pondérées (ou d’absence de mises à jour) appliquées à chaque élément du vecteur d’entrée initial.
  
  
• Enfin, ce vecteur de mises à jour est additionné à la copie du vecteur d’entrée initial qui avait été envoyée directement en sortie via la connexion résiduelle.
  
  
• Le vecteur d’entrée initial est ainsi mis à jour pour refléter la compréhension contextuelle fournie par le SSM sélectif. Il peut alors être transmis à la couche Mamba suivante ou, dans les couches finales du modèle, servir d’entrée à une fonction softmax calculant la probabilité respective que le vecteur entièrement mis à jour corresponde à chaque mot du vocabulaire du modèle.

Un an après l’article original, Tri Dao et Albert Gu ont publié « Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality . Ce document apporte trois contributions majeures :

• Une exploration des liens théoriques entre Mamba et les transformers et un vocabulaire unifié entre les deux architectures
• Une série de clarifications et d’analyses concernant divers choix de conception pour les modèles Mamba
• Une architecture modifiée, Mamba-2, informée et améliorée par ces explorations

L’algorithme Mamba-2 est nettement plus rapide et plus simple à implémenter que le premier Mamba : les auteurs ont fourni un code « minimal SSD » qui implémente un SSM sélectif en 25 lignes de code environ.² Cette efficacité permet à Mamba-2 d’utiliser des dimensions d’état caché beaucoup plus grandes sans ralentir le modèle, ouvrant la voie à des modèles plus grands, plus puissants et plus expressifs basés sur cette architecture. Lors des tests, les modèles Mamba-2 ont égalé ou surpassé les modèles Mamba et les transformers de taille équivalente sur une série de tâches en aval.

Comme le précisent Tri Dao et Albert Gu dans l’introduction de leur article, leur objectif principal était de développer un riche corpus de liens théoriques entre les SSM structurés et les variantes de l’attention. Cela a conduit à un nouveau cadre conceptuel unifiant ces deux aspects, appelé « dualité d’espace d’état (SSD).³ Ce cadre a permis à Mamba de bénéficier de plusieurs années de recherches et d’optimisations réalisées sur l’architecture des transformers.

L’un des avantages notables a été le développement d’un équivalent Mamba de l’attention multitête (MHA) : un bloc Mamba peut être divisé en plusieurs « têtes Mamba » semblables aux multiples « têtes d’attention » des transformers. Une variante de cette approche, qu’ils estimaient analogue à l’attention de requête groupée, permet d’optimiser davantage l’efficacité grâce au parallélisme des tenseurs dans les GPU.

Dans le bloc Mamba-2, qu’ils appellent le bloc Mamba parallèle (par opposition au bloc « séquentiel » d’origine), les paramètres dépendants de l’entrée ∆, B et C sont générés en parallèle au niveau de la couche de projection initiale. B et C, en particulier, sont dérivées en copiant simplement des parties de x_proj, plutôt qu’en multipliant x_proj à travers des couches linéaires dédiées. En plus de simplifier et de réduire le nombre total de paramètres du modèle, ce parallélisme permet un entraînement à grande échelle nettement plus efficace.⁴

Mamba et les transformers ont chacun leurs forces propres, mais les modèles basés sur Mamba sont généralement supérieurs pour tout ce qui touche à l’utilisation de la mémoire et à la vitesse : selon l’article, Mamba offre un débit 5 fois supérieur à celui des transformers équivalents.

Les transformers, eux, restent extrêmement précis et polyvalents, mais aussi très coûteux en ressources de calcul. Pendant le pré-entraînement (et le réglage fin), les besoins en mémoire du mécanisme d’auto-attention augmentent de manière quadratique avec la longueur de la séquence : doubler la taille du contexte d’une séquence quadruple la consommation de ressources. C’est ce « goulot d’étranglement quadratique » qui limite de plus en plus la vitesse et la disponibilité de la mémoire à mesure que la fenêtre de contexte grandit. En inférence, les besoins en mémoire évoluent de manière linéaire.

L’utilisation de la mémoire d’un modèle Mamba évolue de manière linéaire uniquement pendant l’entraînement. Plus important encore, son utilisation de la mémoire en inférence est constante : quel que soit le nombre de tokens vus par le modèle, le SSM maintient une représentation de taille fixe de l’historique de ses entrées. Cela ouvre la voie à des fenêtres de contexte théoriquement illimitées, uniquement freinées par les limitations matérielles.

Cela dit, la méthode des transformers, plus gourmande en mémoire et computationnellement redondante, présente aussi ses avantages propres. Par exemple, des recherches ont montré que les transformers restent supérieurs à Mamba et Mamba-2 pour certaines tâches requérant un apprentissage en contexte (comme l’apprentissage few-shot), les tâches de copie ou encore le raisonnement sur de très longs contextes.

Heureusement, les forces respectives des transformers et de Mamba ne s’excluent pas mutuellement. L’article Mamba-2 suggère qu’un modèle hybride pourrait surpasser aussi bien les transformers purs que les SSM purs, une hypothèse validée par des recherches menées par NVIDIA plus tard en 2024.⁵ De façon générale, les modèles hybrides semblent combiner les avantages en termes d’efficacité de Mamba, et la nuance ainsi que les performances d’apprentissage en contexte offertes par le mécanisme d’attention plus gourmand en ressources des transformers.

Pour explorer cette piste, IBM Research a collaboré avec Tri Dao et Albert Gu, ainsi qu’avec Minjia Zhang de l’Université de l’Illinois à Urbana-Champaign (UIUC), sur Bamba et Bamba V2. Bamba a d’ailleurs inspiré de nombreux éléments architecturaux d’IBM Granite 4.0.

La recherche sur les modèles hybrides reste très active, en particulier dans la communauté open source.