State Space Models (SSM)

Un modèle à espace d’états (SSM) transforme une séquence d’entrée en une séquence de sortie au moyen d’un vecteur d’état caché, mis à jour à chaque étape. Sa mémoire conserve donc la même taille, quelle que soit la longueur de la séquence. Cette idée vient de la théorie du contrôle, où un système linéaire s’écrit sous la forme d’une mise à jour d’état h' = Ah + Bx et d’une sortie y = Ch + Dx, et où le filtre de Kalman estime ces états à partir de mesures bruitées depuis 1960. Les SSM d’apprentissage profond apprennent A, B et C à partir des données et empilent de nombreuses couches de ce type. Depuis S4 en 2021 et Mamba en 2023, ils rivalisent avec les transformeurs sur de longues séquences de texte, d’audio et de séries temporelles, tandis que leurs calculs augmentent linéairement avec la longueur de la séquence.

Fonctionnement des SSM#

Une couche SSM neuronale convertit les équations continues en une récurrence discrète avec un pas appris Δ, selon la règle de maintien d’ordre zéro Ā = exp(ΔA). À chaque étape, elle calcule h_t = Ā h_(t-1) + B̄ x_t et y_t = C h_t. Lorsque A, B et C restent fixes d’une étape à l’autre, le déploiement de la récurrence montre que la sortie est une convolution de l’entrée avec le noyau (CB̄, CĀB̄, CĀ²B̄, ...). La même couche peut donc s’exécuter comme une convolution pendant l’entraînement, en traitant toute la séquence en parallèle sur un GPU, et comme une récurrence pendant la génération, avec une mise à jour de l’état de taille constante par token. Avec PyTorch, les deux approches produisent les mêmes sorties :

import torch

torch.manual_seed(0)
L, N = 64, 16  # longueur de séquence, taille de l’état
x = torch.randn(L)

# SSM continu diagonal, discrétisé avec un pas delta (maintien d’ordre zéro)
A = -(torch.rand(N) + 0.5)
B, C, delta = torch.randn(N), torch.randn(N), 0.1
A_bar = torch.exp(delta * A)
B_bar = (A_bar - 1) / A * B

# Approche récurrente : une mise à jour de l’état par étape, comme pendant la génération
h, y_rec = torch.zeros(N), []
for t in range(L):
    h = A_bar * h + B_bar * x[t]
    y_rec.append(C @ h)
y_rec = torch.stack(y_rec)

# Approche convolutionnelle : précalcul du noyau K_k = C . (A_bar^k * B_bar), comme pendant l’entraînement
K = torch.stack([C @ (A_bar**k * B_bar) for k in range(L)])
y_conv = torch.stack([K[: t + 1].flip(0) @ x[: t + 1] for t in range(L)])

print("Same outputs:", torch.allclose(y_rec, y_conv, atol=1e-5))

Le programme affiche Same outputs: True. La récurrence ne conserve que 16 nombres d’état à chaque étape, tandis que le noyau précalculé permet à l’entraînement de calculer les 64 sorties en parallèle.

De S4 à Mamba#

Les premiers SSM profonds étaient trop lents et trop gourmands en mémoire pour être entraînés à grande échelle. S4, conçu par Albert Gu, Karan Goel et Christopher Ré, a structuré la matrice d’état de façon à rendre le calcul du noyau de convolution peu coûteux, et a introduit une initialisation conçue pour mémoriser de longs historiques. Il a établi de nouvelles performances de pointe pour toutes les tâches du benchmark Long Range Arena, est devenu le premier modèle à résoudre la tâche Path-X sur des séquences de 16 000 étapes et a généré 60 fois plus vite qu’un transformeur.

Un noyau fixe traite toutes les entrées de la même façon ; il ne peut donc pas décider d’ignorer du contenu superflu et de conserver un fait important. Mamba, conçu par Albert Gu et Tri Dao, fait de Δ, B et C des fonctions du token courant, ce qui permet au modèle de choisir ce qu’il écrit dans son état et ce qu’il oublie. Ces paramètres sélectifs empêchent d’utiliser l’approche par convolution ; Mamba calcule donc la récurrence avec un parcours parallèle adapté au matériel. Mamba a atteint un débit d’inférence cinq fois supérieur à celui des transformeurs, et Mamba-3B a surpassé des transformeurs de même taille tout en égalant des transformeurs deux fois plus grands. Chaque bloc Mamba entoure le SSM sélectif d’une courte convolution et d’une porte SiLU :

Mamba-2 a montré que les SSM sélectifs et une forme d’attention linéaire sont deux représentations d’un même calcul, dans un cadre que ses auteurs appellent la dualité espace d’états. Cette approche a permis de rendre la couche centrale 2 à 8 fois plus rapide.

SSM et architectures apparentées#

• Transformeurs : l’attention automatique compare chaque token à tous les autres, ce qui lui permet de copier exactement n’importe quel token précédent, au prix d’un coût qui augmente de manière quadratique et d’un cache KV dont la taille croît avec la fenêtre de contexte. Un SSM compresse l’historique dans un état fixe : c’est moins coûteux, mais avec perte. L’article de 2024 Repeat After Me a montré que les transformeurs copient et récupèrent mieux des informations dans de longues entrées que les SSM à état fixe.
• Réseaux neuronaux récurrents (RNNs) : un RNN ou un LSTM applique une mise à jour non linéaire, si bien que l’entraînement doit parcourir la séquence token par token. Les mises à jour des SSM sont linéaires par rapport à l’état, ce qui permet d’utiliser une convolution ou un parcours parallèle.
• Modèles de Markov cachés (HMMs) : un HMM comporte un ensemble fini d’états discrets et des probabilités de transition. Le modèle linéaire gaussien à espace d’états qui sous-tend le filtre de Kalman en est l’équivalent continu, et les SSM profonds conservent ce vecteur d’état continu.

Applications des SSM#

La plupart des usages en production concernent des grands modèles de langage (LLMs) hybrides qui conservent quelques couches d’attention pour retrouver précisément les informations. Le Jamba d’AI21 alterne des couches de transformeur et de Mamba avec des couches à mélange d’experts (MoE), tient sur un seul GPU de 80 Go et affiche de solides résultats avec des contextes allant jusqu’à 256 000 tokens. Les modèles hybrides Granite 4.0-H d’IBM combinent des couches Mamba-2 et des couches de transformeur selon un ratio de 9:1. Vision Mamba applique des parcours sélectifs aux patches d’image, et les SSM modélisent aussi des séquences brutes d’audio et d’ADN trop longues pour une attention complète. La forme classique est utilisée dans le suivi avec Ultralytics YOLO26 : le traqueur TrackTrack par défaut, comme ByteTrack et BoT-SORT, prédit la prochaine boîte de chaque objet avec un filtre de Kalman dont l’état à 8 dimensions contient le centre et la taille de la boîte ainsi que leurs vitesses. Cet état est mis à jour à l’aide d’une matrice de transition à vitesse constante et d’une matrice d’observation linéaire qui lit les quatre coordonnées de la boîte.