State Space Models (SSM)

Ein Zustandsraummodell (SSM) bildet eine Eingabesequenz über einen verborgenen Zustandsvektor, den es bei jedem Schritt aktualisiert, auf eine Ausgabesequenz ab. Dadurch bleibt sein Speicherbedarf unabhängig von der Sequenzlänge gleich. Die Idee stammt aus der Regelungstechnik. Dort wird ein lineares System durch eine Zustandsaktualisierung h' = Ah + Bx und eine Ausgabe y = Ch + Dx beschrieben. Der Kalman-Filter schätzt solche Zustände seit 1960 aus verrauschten Messungen. SSMs im Deep Learning lernen A, B und C anhand von Daten und stapeln viele dieser Schichten. Seit S4 im Jahr 2021 und Mamba im Jahr 2023 konkurrieren sie bei langen Sequenzen aus Text-, Audio- und Zeitreihendaten mit Transformern, wobei ihr Rechenaufwand linear mit der Sequenzlänge wächst.

So funktionieren SSMs#

Eine neuronale SSM-Schicht wandelt die kontinuierlichen Gleichungen mit einer erlernten Schrittweite Δ in eine diskrete Rekursion um. Dabei kommt die Zero-Order-Hold-Regel Ā = exp(ΔA) zum Einsatz. Bei jedem Schritt werden dann h_t = Ā h_(t-1) + B̄ x_t und y_t = C h_t berechnet. Bleiben A, B und C über die Zeitschritte hinweg konstant, zeigt das Aufrollen der Rekursion, dass die Ausgabe eine Faltung der Eingabe mit dem Kernel (CB̄, CĀB̄, CĀ²B̄, ...) ist. Dieselbe Schicht kann daher während des Trainings als Faltung ausgeführt werden, wobei die gesamte Sequenz parallel auf einer GPU verarbeitet wird. Bei der Generierung lässt sie sich dagegen als Rekursion ausführen, mit einer Aktualisierung des Zustands konstanter Größe pro Token. In PyTorch liefern beide Ansichten dieselben Ausgaben:

import torch

torch.manual_seed(0)
L, N = 64, 16  # Sequenzlänge, Zustandsgröße
x = torch.randn(L)

# Diagonales SSM mit kontinuierlicher Zeit, diskretisiert mit der Schrittweite delta (Zero-Order Hold)
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

# Rekursive Ansicht: eine Zustandsaktualisierung pro Schritt, wie bei der Generierung
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)

# Faltungsansicht: Kernel K_k = C . (A_bar^k * B_bar) vorberechnen, wie beim Training
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))

Die Ausgabe lautet Same outputs: True. Die Rekursion behält bei jedem Schritt nur 16 Zustandswerte bei. Mit dem vorberechneten Kernel kann das Training alle 64 Ausgaben parallel berechnen.

Von S4 zu Mamba#

Die ersten tiefen SSMs waren für ein Training im großen Maßstab zu langsam und speicherhungrig. S4 von Albert Gu, Karan Goel und Christopher Ré gab der Zustandsmatrix eine Struktur, mit der sich der Faltungskernel kostengünstig berechnen lässt, und eine Initialisierung, die auf das Behalten langer Verläufe ausgelegt ist. S4 setzte bei jeder Aufgabe des Long Range Arena-Benchmarks einen neuen Spitzenwert, löste als erstes Modell die Path-X-Aufgabe mit Sequenzen von 16.000 Schritten und generierte 60-mal schneller als ein Transformer.

Ein fester Kernel behandelt jede Eingabe gleich. Daher kann er nicht entscheiden, Füllmaterial zu überspringen und eine wichtige Information zu behalten. Mamba von Albert Gu und Tri Dao macht Δ, B und C von Funktionen des aktuellen Tokens abhängig. So kann das Modell entscheiden, was es in seinen Zustand schreibt und was es vergisst. Diese selektiven Parameter verhindern die Darstellung als Faltung, weshalb Mamba die Rekursion mit einem hardwarebewussten parallelen Scan berechnet. Mamba erreichte den fünffachen Inferenzdurchsatz von Transformern. Mamba-3B übertraf Transformer gleicher Größe und erreichte das Niveau von Transformern mit doppelt so vielen Parametern. Jeder Mamba-Block ergänzt das selektive SSM um eine kurze Faltung und ein SiLU-Gate:

Mamba-2 zeigte, dass selektive SSMs und eine Form von linearer Aufmerksamkeit zwei Ansichten derselben Berechnung sind. Die Autoren nennen diesen Ansatz Zustandsraum-Dualität und nutzten ihn, um die Kernschicht zwei- bis achtmal schneller zu machen.

SSMs im Vergleich zu verwandten Architekturen#

• Transformer: Selbstaufmerksamkeit vergleicht jedes Token mit jedem anderen. So kann sie jedes frühere Token exakt kopieren. Dafür wächst der Rechenaufwand quadratisch und der KV-Cache mit dem Kontextfenster. Ein SSM komprimiert den Verlauf in einen Zustand fester Größe. Das ist kostengünstiger, geht aber mit Informationsverlust einher. Die Arbeit Repeat After Me aus dem Jahr 2024 zeigte, dass Transformer lange Eingaben besser kopieren und Informationen daraus abrufen als SSMs mit festem Zustand.
• Rekurrente neuronale Netze (RNNs): Ein RNN oder LSTM wendet eine nichtlineare Aktualisierung an. Daher muss das Training die Sequenz Token für Token durchlaufen. SSM-Aktualisierungen sind linear in Bezug auf den Zustand, was eine Faltung oder einen parallelen Scan ermöglicht.
• Verdeckte Markov-Modelle (HMMs): Ein HMM hat eine endliche Menge diskreter Zustände mit Übergangswahrscheinlichkeiten. Das lineare gaußsche Zustandsraummodell hinter dem Kalman-Filter ist das kontinuierliche Gegenstück dazu. Tiefe SSMs behalten diesen kontinuierlichen Zustandsvektor bei.

Einsatzbereiche von SSMs#

Die meisten Anwendungen im Produktiveinsatz finden sich in hybriden großen Sprachmodellen (LLMs), die einige Aufmerksamkeitsschichten für den exakten Abruf beibehalten. Jamba von AI21 verschachtelt Transformer- und Mamba-Schichten mit Schichten für eine Expertenmischung (MoE), passt auf eine einzelne GPU mit 80 GB und erzielt laut Angaben bei Kontexten mit bis zu 256.000 Token gute Ergebnisse. Die hybriden Granite-4.0-H-Modelle von IBM kombinieren Mamba-2- und Transformer-Schichten im Verhältnis 9:1. Vision Mamba wendet selektive Scans auf Bildpatches an. SSMs modellieren außerdem rohe Audio- und DNA-Sequenzen, die für vollständige Aufmerksamkeit zu lang sind. Die klassische Form kommt beim Tracking mit Ultralytics YOLO26 zum Einsatz: Der standardmäßige TrackTrack-Tracker sagt wie ByteTrack und BoT-SORT die nächste Bounding Box jedes Objekts mit einem Kalman-Filter voraus. Dessen achtdimensionaler Zustand enthält die Mitte und Größe der Bounding Box sowie deren Geschwindigkeiten. Er wird mit einer Übergangsmatrix für konstante Geschwindigkeit und einer linearen Beobachtungsmatrix aktualisiert, die die vier Koordinaten der Bounding Box ausliest.