Mamba-3 : un modèle à espace d'états conçu pour l'inférence

Cet article est repris de Goomba Lab, dirigé par Albert Gu, responsable scientifique de Cartesia.

Depuis la sortie de Mamba-2 à la mi-2024, la plupart des architectures ont abandonné Mamba-1. Pourquoi ? Mamba-2 faisait le pari que l’efficacité de l’entraînement était le principal frein des modèles à espace d’états, ou SSM. Il a donc simplifié leur mécanisme pour entraîner deux à huit fois plus vite que son prédécesseur, favorisant une adoption plus large.

Depuis, le monde des LLM évolue. Le préentraînement reste très important, mais l’attention se porte davantage sur le post-entraînement et le déploiement, tous deux très gourmands en inférence. Le développement des méthodes de post-entraînement, notamment l’apprentissage par renforcement avec récompenses vérifiables, ou RLVR, pour le code et les mathématiques, exige de générer énormément de trajectoires. Plus récemment, les processus agentiques comme Codex, Claude Code ou OpenClaw ont fait exploser la demande d’inférence.

Malgré cette importance croissante, de nombreuses architectures linéaires, dont Mamba-2, ont été conçues d’abord pour l’entraînement. Pour accélérer le préentraînement, le SSM a été progressivement simplifié, par exemple en réduisant la transition diagonale à un scalaire multiplié par l’identité. L’entraînement y a gagné en vitesse, mais l’étape d’inférence est devenue trop simple et limitée par la mémoire. Les GPU passent surtout leur temps à déplacer des données plutôt qu’à calculer.

À cette nouvelle époque de l’inférence, nous voulons repousser la frontière entre qualité et efficacité : faire fonctionner les meilleurs modèles plus vite.

Une question se pose naturellement :

À quoi ressemblerait un SSM conçu pour l’inférence ?

Le modèle Mamba-3

Que manque-t-il ? L’intérêt des modèles linéaires est dans leur nom : le calcul augmente linéairement avec la longueur des séquences grâce à un état fixe. Mais ce gain a une contrepartie. La taille d’état fixe impose de compresser toute l’information passée en une représentation, à l’opposé d’un Transformer qui la conserve dans un état croissant, le cache KV. Si l’on ne peut pas agrandir l’état, comment lui faire faire davantage de travail ?

Les conceptions précédentes simplifiaient la récurrence et la matrice de transition pour accélérer l’entraînement. Cela réduisait aussi la richesse de la dynamique et laissait le décodage limité par la mémoire : chaque token exigeait peu de calcul par rapport aux données déplacées. Nous pouvons donc agir de trois façons : rendre la récurrence plus expressive, utiliser une matrice de transition plus riche et ajouter du travail parallèle presque gratuit à chaque mise à jour.

À partir de ces observations, nous améliorons Mamba-2 de trois façons :

1. Nous augmentons l’expressivité du mécanisme SSM avec une récurrence plus générale issue de notre schéma de discrétisation exponentiel-trapézoïdal.
2. Nous étendons les capacités de suivi d’état en modélisant un système SSM à valeurs complexes.
3. Nous améliorons les performances générales avec peu d’effet sur la latence de décodage grâce aux SSM à entrées et sorties multiples, ou MIMO, qui modélisent plusieurs SSM en parallèle au lieu des SSM actuels à entrée et sortie uniques, ou SISO.

Grâce à ces changements, Mamba-3 repousse les performances tout en conservant une latence d’inférence similaire.

Ces trois changements s’inspirent de la littérature plus classique sur la théorie du contrôle et les modèles à espace d’états.

Nos travaux vont à contre-courant de nombreuses architectures linéaires modernes qui utilisent d’autres interprétations de la récurrence, comme l’attention linéaire ou l’entraînement au moment du test, qui représentent difficilement ces concepts.

Architecture

Qu’est-ce qui change dans la couche Mamba-2 ? Outre les trois améliorations du SSM décrites ci-dessus, nous avons remanié l’architecture pour la rapprocher des modèles de langage modernes habituels.

Le schéma montre plusieurs changements. Voici les principaux.

Normalisations. Nous avons ajouté QKNorm 1 ¹, qui stabilise empiriquement l’entraînement de Mamba-3. Cela le rapproche des Transformers contemporains et de Gated DeltaNet, ou GDN. Avec QKNorm, le RMSNorm de Mamba-2 devient facultatif. Nos expériences montrent néanmoins qu’il peut rester utile dans les modèles hybrides en aidant l’extrapolation en longueur. Nous y reviendrons.

Adieu à la convolution courte. Nous avons supprimé la courte convolution causale de Mamba-1/2 en combinant des biais simples sur B et C après BCNorm avec notre nouvelle récurrence issue de la discrétisation. Cette récurrence applique implicitement une convolution à l’entrée de l’état caché. La deuxième partie explique comment.

Deux nouveaux composants apparaissent aussi : RoPE et les projections MIMO. RoPE exprime les SSM complexes en interprétant les transitions complexes comme des rotations, évitant une réimplémentation coûteuse des noyaux. Les projections MIMO étendent les matrices B et C à la représentation nécessaire aux SSM MIMO.

La deuxième partie détaille les raisons et l’implémentation de ces deux composants. Pour l’instant, voyez-les comme des améliorations fondamentales indépendantes qui contribuent chacune aux performances ou aux capacités du modèle.

Enfin, notre architecture adopte des couches MLP intercalées, selon la convention des Transformers et d’autres modèles linéaires.

Résultats empiriques

Nous évaluons notre modèle Mamba-3 final face à d’autres architectures linéaires courantes et à un Transformer de référence.

Modélisation du langage

Mamba-3 dépasse Mamba-2 et des alternatives performantes d’attention linéaire comme GDN en modélisation du langage, à différentes tailles de modèles préentraînés. Mamba-3-SISO est directement comparable aux modèles linéaires précédents. Il a notamment exactement les mêmes dimensions architecturales que Mamba-2, taille du modèle, de l’état, etc., et un temps d’entraînement comparable. Notre variante MIMO améliore encore l’exactitude de plus d’un point de pourcentage sur les tâches en aval à l’échelle d’un milliard de paramètres. Elle exige toutefois un entraînement plus long, mais pas une latence de décodage plus élevée.

Tâches de récupération

Avec leur état fixe, les modèles linéaires sont naturellement moins performants que les Transformers sur les tâches de récupération. Comme prévu, parmi les modèles purs, le Transformer domine ces tâches, mais Mamba-3 se comporte bien parmi les alternatives sous-quadratiques. L’ajout de MIMO améliore encore la récupération sans augmenter la taille de l’état.

Compte tenu de cette faiblesse intrinsèque, mais de bonnes performances générales de modélisation,

nous prévoyons que les couches linéaires seront surtout utilisées avec des couches d’auto-attention globale à l’avenir.*

*au moins pour la modélisation du langage

Les modèles hybrides combinent le fonctionnement général proche d’une mémoire des couches linéaires avec le stockage exact, proche d’une base de données, du cache KV de l’auto-attention. Ils dépassent empiriquement les modèles purs tout en économisant fortement mémoire et calcul ¹⁰. Ici aussi, la combinaison de couches linéaires et d’auto-attention améliore la récupération par rapport à un Transformer classique.

Cependant, l’interaction exacte entre ces modèles linéaires et l’auto-attention n’est pas entièrement comprise. La projection facultative avant sortie de Mamba-3 améliore, par exemple, la généralisation en longueur sur les tâches synthétiques NIAH, au prix d’une légère baisse sur les tâches réelles de récupération en contexte. Même les détails de la normalisation réintroduite, sa position avant ou après la porte et son type groupé ou classique, influencent sensiblement l’exactitude sur des données semi-structurées et non structurées comme FDA et SWDE.

Des noyaux partout

Nous avons hâte de voir ce que vous construirez avec Mamba-3. Pour vous y aider, nous publions nos noyaux en open source. Ils sont aussi rapides que les noyaux Triton originaux de Mamba-2.

Mesurer les latences

Latence de préremplissage

Latence de préremplissage et décodage

À l’échelle de 1,5 milliard de paramètres, Mamba-3 SISO atteint la plus faible latence cumulée de préremplissage et de décodage pour toutes les longueurs de séquence. Il dépasse Mamba-2, Gated DeltaNet et même le Transformer avec son écosystème vLLM très optimisé. Mamba-3 MIMO a une vitesse comparable à Mamba-2, avec des performances nettement supérieures.

Le préremplissage Triton de Mamba-3 SISO conserve des performances presque identiques à Mamba-2. La nouvelle discrétisation et les plongements RoPE dépendant des données n’ajoutent donc pas de surcoût. Mamba-3 MIMO ne subit qu’un ralentissement modéré du préremplissage grâce à son implémentation TileLang efficace. Les bons résultats de décodage des deux variantes viennent en partie de l’implémentation CuTe DSL, facilitée par la simplicité des composants de Mamba-3.

Choix de conception

Nous avons beaucoup réfléchi à la façon d’accélérer au maximum les noyaux sans compliquer leur utilisation. Nous avons retenu Triton, TileLang et CuTe DSL.

Le choix de Triton était assez simple. Il est presque standard pour développer des architectures, comme le montre le dépôt flash linear attention entièrement en PyTorch et Triton. Il dépasse PyTorch standard grâce au contrôle du découpage en tuiles et à la fusion de noyaux, tout en restant indépendant de la plateforme. Il offre aussi l’injection de PTX, un langage assembleur pour GPU, et le Tensor Memory Accelerator des GPU Hopper pour des transferts asynchrones massifs de la mémoire globale vers la mémoire partagée.

Nous avons développé les noyaux de préremplissage MIMO avec TileLang. Les projections supplémentaires permettent de réduire les entrées-sorties mémoire en exploitant la hiérarchie mémoire du GPU. Triton n’offrait pas le contrôle fin souhaité. TileLang permet de déclarer et de contrôler explicitement les tuiles en mémoire partagée et de créer des fragments de registres. Nous réutilisons ainsi mieux la mémoire tout en gardant un niveau d’abstraction assez élevé pour développer rapidement.

Étant donné l’importance de l’inférence et du décodage, nous avons choisi CuTe DSL pour nos noyaux de décodage. Son interface Python génère des noyaux bas niveau à partir des abstractions de CUTLASS. Nous disposons pratiquement du contrôle de CUDA pour créer des noyaux performants adaptés à notre matériel, ici les GPU Hopper. Le contrôle fin de l’organisation des tenseurs et de la spécialisation des warps nous permet d’exploiter toutes les capacités du GPU.

Ces implémentations à différents niveaux d’abstraction GPU sont rendues possibles par la conception algorithmique des ajouts simples et légers de Mamba-3 et par leur mise en œuvre. Notre publication complète détaille notamment la structure de fusion et le langage spécialisé des noyaux.

La suite

Merci d’avoir lu cette première partie ! Nous n’avons pas pu couvrir tous les détails des noyaux, résultats expérimentaux et études d’ablation. Vous les trouverez dans notre article, et les noyaux sont publiés dans mamba-ssm.

La deuxième et dernière partie examine les trois améliorations centrales de Mamba-3, leurs fondements SSM et quelques pistes qui nous intéressent particulièrement.

Notes