Mamba expliqué : les Modèles d’Espace d’État promettent des performances proches des Transformers avec des contextes d’un million de tokens

TL;DR

• Mamba utilise un Modèle d’Espace d’État (SSM) inspiré de la théorie du contrôle pour la communication, remplaçant le mécanisme Attention dans les Transformers.
• Il promet des performances similaires et des lois de scaling proches de celles des Transformers tout en permettant des contextes longs jusqu’à environ un million de tokens.
• Dans les benchmarks, Mamba-3B dépasse les Transformers de même taille et égalise des Transformers deux fois plus gros lors du pré-entraînement et des évaluations en aval, avec une inférence rapide (rapportée jusqu’à 5x plus rapide que le Transformer).
• L’architecture conserve des projections de type MLP pour le calcul, et des convolutions locales; le SSM gère la communication entre les tokens. Mamba Explained on The Gradient

Contexte et arrière-plan

Les Transformers dominent actuellement les avancées majeures en IA, mais présentent des inefficacités pour les séquences longues. Le problème central est l’Attention, où chaque token peut regarder tous les tokens précédents, créant un goulot d’étranglement quadratique en entraînement (O(n^2)) et nécessitant une mémoire O(n) pour le KV cache lors de la génération autoregressive. Bien que des techniques comme l’Attention à fenêtre glissante ou des optimisations comme FlashAttention puissent atténuer certains coûts, elles ne résolvent pas fondamentalement le besoin d’une approche différente pour de très longs contextes. Parallèlement, les Modèles d’Espace d’État (SSM) proposent une perspective différente. Les auteurs de Mamba positionnent les SSM comme une base viable pour le traitement des séquences, tout en conservant des projections MLP pour le calcul et en utilisant les SSM pour la communication entre les tokens. L’idée est de substituer la composante Attention par un SSM qui capture la dynamique temporelle de manière compacte. L’article motive cela avec un exemple type Temple Run pour illustrer comment un système peut inférer des actions futures à partir d’un ensemble limité d’observations en exploité la structure de l’évolution de l’état. L’objectif est d’atteindre un équilibre entre efficacité et performance comparable à la frontière de Pareto. Mamba Explained on The Gradient Une étape conceptuelle clé consiste à modéliser le temps comme continu et à le convertir ensuite en temps discret pour une utilisation pratique. Mamba utilise une discrétisation par Zero-Order Hold (ZOH) pour transformer l’équation d’état continue en une équation en temps discret. L’article montre comment l’état caché h et l’entrée x influencent la prochaine action et la sortie, via des matrices A, B, C, D et un pas de temps ∆ qui détermine l’influence du passé. Cette perspective positionne Mamba comme un remplacement direct de l’attention au sein du bloc Mamba, tout en conservant la computation via des projections linéaires et des convolutions locales. Mamba Explained on The Gradient La discussion oppose ensuite la transmission d’information des Transformers et l’approche SSM, à travers l’analogie du problème de la fête du cocktail: l’attention dans les Transformers peut être vue comme la mémoire photographique du passé, alors que le SSM offre un profil d’efficacité différent, plus proche des RNNs classiques en termes d’utilisation de mémoire, mais avec des capacités de modélisation de séquence accrues. Le résultat est une voie proposée vers une frontière de Pareto plus large entre efficacité et efficacité. Mamba Explained on The Gradient

Ce qui est nouveau

L’élément central de Mamba est de remplacer le mécanisme de communication basé sur l’attention dans un bloc Transformer par un Modèle d’Espace d’État (SSM) pour la communication, tout en conservant des projections de calcul de type MLP et des convolutions locales. Le bloc Mamba offre une architecture en blocs similaire à celle des Transformer blocs, mais avec une communication entre tokens gérée par le SSM, permettant des contextes plus longs avec un escalonement linéaire. Les éléments clés incluent :

• Formulation continuelle du temps: h’(t) = A h(t) + B x(t); y(t) = C h(t) + D x(t).
• Discrétisation: h_{t+1} ≈ (I + ∆ A) h_t + (∆ B) x_t; y_t ≈ C h_t + D x_t; ∆ est le pas de temps, qui régule l’influence du passé sur le prochain état.
• Le SSM est présenté comme un substitut direct à la composante de communication dans chaque bloc Mamba, permettant un traitement sur longue distance sans le goulot d’étranglement quadratique de l’attention.
• Le chemin computationnel repose sur des projections linéaires, des non-linéarités et des convolutions locales, préservant une base ML familière tout en modifiant la communication entre tokens. | Caractéristique | Transformer | Mamba |--- |--- |--- |Mécanisme de communication | Attention avec KV cache | Modèle d’Espace d’État (SSM) pour la communication |Complexité temps/espace | Entraînement O(n^2); génération O(n) par token; mémoire O(n) | Escalade linéaire en fonction de la longueur de la séquence |Bottleneck computationnel | Attention + MLPs | Projections de type MLP + convolutions locales |Scalabilité du contexte | Quadratique avec le contexte | Linéaire avec le contexte; conçu pour contextes longs (jusqu’à ~1M tokens) |

Points clés

• Mamba introduit des Modèles d’Espace d’État comme mécanisme de communication pour remplacer l’attention dans les Transformers, visant des contextes longs avec un escalade linéaire.
• La voie computationnelle reste basée sur des projections de type MLP et des convolutions locales, mais la communication entre tokens est gérée par le SSM.
• La discrétisation par Zero-Order Hold traduit la dynamique en continu en temps discret, avec ∆ contrôlant l’influence du passé.
• Des résultats rapportés positionnent le Mamba-3B comme équivalent ou meilleur que des Transformers de même taille et compétitif avec des Transformers plus gros, avec une inférence plus rapide dans la pratique.
• Si ces résultats se vérifient largement, les blocs de style Mamba pourraient influencer la façon dont les entreprises abordent le calcul de contextes longs et le déploiement efficace.

Références

• https://thegradient.pub/mamba-explained