Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu; Tri Dao

Paper · arXiv 2312.00752

Mamba: Lineare Sequenzmodellierung mit selektiven Zustandsräumen

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu, Tri Dao

Veröffentlicht 2023-12-01cs.LGcs.AI arXiv PDF

Mit Übersetzung und KI lesen

Zusammenfassung in drei Zeilen

🎯 Transformer skalieren quadratisch mit der Sequenzlänge und erfordern während der autoregressiven Inferenz große KV-Caches, während frühere subquadratische Modelle wie LTI-Zustandsraummodelle bei diskreten, inhaltsbasierten Aufgaben wie der Sprachmodellierung an ihre Grenzen stoßen.
🔧 Macht SSM-Parameter eingabeabhängig (selektiv), um Informationen dynamisch zu filtern oder weiterzuleiten, und kombiniert dies mit einem hardware-optimierten, fusionierten parallelen assoziativen Scan sowie einem Recomputation-Algorithmus, um die Materialisierung des expandierten verborgenen Zustands im GPU-HBM zu vermeiden.
📈 Mamba-3B erreicht oder übertrifft doppelt so große Transformer bei nachgelagerten Zero-Shot-Aufgaben (z. B. 63,3 % Durchschnitt gegenüber 59,1 % bei Pythia-2.8B und 61,7 % bei Pythia-6.9B) und bietet gleichzeitig einen bis zu 5-mal höheren Inferenzdurchsatz.

Abstract

Basismodelle, die heute die meisten der spannendsten Anwendungen im Deep Learning antreiben, basieren fast ausnahmslos auf der Transformer-Architektur und ihrem zentralen Attention-Modul. Viele Architekturen mit subquadratischer Laufzeit wie lineare Attention, gaterte Faltungs- und rekurrente Modelle sowie strukturierte Zustandsraummodelle (SSMs) wurden entwickelt, um die rechnerische Ineffizienz von Transformern bei langen Sequenzen zu beheben; auf wichtigen Modalitäten wie Sprache erreichten sie jedoch nicht die Qualität von Attention. Wir identifizieren das Unvermögen zu inhaltsbasiertem logischen Schließen als zentrale Schwachstelle solcher Modelle und führen mehrere Verbesserungen ein. Erstens behebt das bloße Definieren der SSM-Parameter als Funktionen der Eingabe deren Schwäche bei diskreten Modalitäten, da das Modell Informationen entlang der Sequenzlängendimension in Abhängigkeit vom aktuellen Token selektiv weiterleiten oder verwerfen kann. Zweitens entwerfen wir einen hardware-optimierten parallelen Algorithmus im rekurrenten Modus, obwohl diese Änderung die Verwendung effizienter Faltungen verhindert. Wir integrieren diese selektiven SSMs in eine vereinfachte neuronale End-to-End-Netzwerkarchitektur ohne Attention- oder gar MLP-Blöcke (Mamba). Mamba profitiert von schneller Inferenz (5 $\times$ höherer Durchsatz als Transformer) sowie linearer Skalierung bezüglich der Sequenzlänge, und seine Leistung verbessert sich auf realen Daten bis hin zu Sequenzen mit einer Million Tokens. Als universelles Sequenzmodell-Backbone erzielt Mamba Spitzenleistungen über verschiedene Modalitäten wie Sprache, Audio und Genomik hinweg. Bei der Sprachmodellierung übertrifft unser Mamba-3B-Modell gleich große Transformer und erreicht sowohl beim Vortraining als auch in der Downstream-Evaluation die Qualität doppelt so großer Transformer.

Original-Abstract (Englisch)

Foundation models, now powering most of the exciting applications in deep learning, are almost universally based on the Transformer architecture and its core attention module. Many subquadratic-time architectures such as linear attention, gated convolution and recurrent models, and structured state space models (SSMs) have been developed to address Transformers' computational inefficiency on long sequences, but they have not performed as well as attention on important modalities such as language. We identify that a key weakness of such models is their inability to perform content-based reasoning, and make several improvements. First, simply letting the SSM parameters be functions of the input addresses their weakness with discrete modalities, allowing the model to selectively propagate or forget information along the sequence length dimension depending on the current token. Second, even though this change prevents the use of efficient convolutions, we design a hardware-aware parallel algorithm in recurrent mode. We integrate these selective SSMs into a simplified end-to-end neural network architecture without attention or even MLP blocks (Mamba). Mamba enjoys fast inference (5 $\times$ higher throughput than Transformers) and linear scaling in sequence length, and its performance improves on real data up to million-length sequences. As a general sequence model backbone, Mamba achieves state-of-the-art performance across several modalities such as language, audio, and genomics. On language modeling, our Mamba-3B model outperforms Transformers of the same size and matches Transformers twice its size, both in pretraining and downstream evaluation.

Leseleitfaden

Problemstellung und Motivation

Basismodelle stützen sich überwiegend auf die Transformer-Architektur und Self-Attention. Obwohl dichte Attention Informationen effektiv über ein Sequenz-Kontextfenster hinweg weiterleitet, birgt sie zwei fundamentale Einschränkungen: quadratische Rechenkosten $O (L^{2})$ in Bezug auf die Sequenzlänge $L$ und die Notwendigkeit, vergangene Keys und Values zu speichern (der KV-Cache), was zu linearer Inferenzzeit pro Schritt und Speicherengpässen führt.

Bisherige subquadratische Ansätze – wie lineare Attention, rekurrente Architekturen und strukturierte Zustandsraummodelle (SSMs) wie S4 – skalieren zeitlich linear oder pseudolinear. Historisch gesehen konnten strukturierte SSMs jedoch bei diskreten, informationsdichten Modalitäten wie Text nicht mit der Qualität von Transformern mithalten. Die Autoren identifizieren die Ursache: Bisherige SSMs sind zeitinvariant (LTI, linear time-invariant). Da sich ihre Übergangsdynamik nicht in Abhängigkeit vom Inhalt der Eingabe ändert, können sie den Kontext nicht selektiv komprimieren, irrelevantes Rauschen herausfiltern oder die Informationsspeicherung über variable Abstände hinweg modulieren.

Zentrale Ideen

Selektionsmechanismus als Kontextkompression: Sequenzmodellierung erfordert das Komprimieren der Historie in einen Zustand. Frühere LTI-SSMs nutzen feste Übergangsparameter, wodurch sie alle Zeitschritte einheitlich behandeln müssen. Indem die SSM-Matrizen auf Basis des aktuellen Eingabe-Tokens parametrisiert werden, erlangt das Modell das erforderliche Inhaltsverständnis, um Informationen selektiv zu filtern oder zu behalten.
Überwindung des Faltungsengpasses: Werden Parameter eingabeabhängig gemacht, bricht die mathematische Äquivalenz zwischen linearen Rekurrenzen und globalen Faltungen zusammen. Ohne Faltungen würde eine naive Rekurrenz erfordern, einen expandierten latenten Zustand der Größe $(B, L, D, N)$ im GPU-High-Bandwidth-Memory (HBM) zu materialisieren. Mamba löst dies, indem Diskretisierung, paralleler assoziativer Scan und Ausgabeprojektion im schnellen On-Chip-SRAM der GPU fusioniert werden, wodurch der HBM-Transfer minimal bleibt.
Einheitliche Architektur: Mamba ersetzt verschachtelte Multi-Head-Attention- (MHA) und MLP-Blöcke durch einen einzigen homogenen Block, der das Gated-Design von H3 mit dem Gated-MLP im SwiGLU-Stil vereint.

Methode

Kontinuierliche Zustandsraummodelle bilden eine Eingabe $x (t) \in R$ über einen impliziten latenten Zustand $h (t) \in R^{N}$ auf $y (t) \in R$ ab: $h^{'} (t) = A h (t) + B x (t), y (t) = C h (t)$

Unter Verwendung der Zero-Order-Hold-Diskretisierung (ZOH) mit Schrittweite $Δ$ transformiert sich das kontinuierliche System in diskrete Rekurrenzparameter: $\overline{A} = exp (Δ A), \overline{B} = (Δ A)^{- 1} (exp (Δ A) - I) \cdot Δ B$ $h_{t} = \overline{A} h_{t - 1} + \overline{B} x_{t}, y_{t} = C h_{t}$

Der Selektionsmechanismus (S6)

Anstatt $(Δ, B, C)$ über die Zeit konstant zu halten, macht Mamba sie eingabeabhängig:

$B_{t} = s_{B} (x_{t}) = Linear_{N} (x_{t})$
$C_{t} = s_{C} (x_{t}) = Linear_{N} (x_{t})$
$Δ_{t} = τ_{Δ} (Parameter + s_{Δ} (x_{t}))$ , mit $τ_{Δ} = softplus$ und $s_{Δ} (x) = Broadcast_{D} (Linear_{1} (x))$

Dies generalisiert das klassische RNN-Gating: Wenn $N = 1, A = - 1, B = 1$ , vereinfacht sich die Rekurrenz direkt zu einem gegateten Update $h_{t} = (1 - g_{t}) h_{t - 1} + g_{t} x_{t}$ , wobei $g_{t} = σ (Linear (x_{t}))$ .

Hardware-optimierte Zustandsexpansion

Um einen hohen Durchsatz aufrechtzuerhalten, ohne den $(B, L, D, N)$ -Zustand im HBM zu materialisieren:

Parameter $(Δ, A, B, C)$ werden aus dem HBM direkt in das On-Chip-SRAM geladen.
Die Diskretisierung in $(\overline{A}, \overline{B})$ und ein arbeitseffizienter paralleler assoziativer Scan werden vollständig im SRAM berechnet.
Der latente Zustand wird mit $C$ auf die Dimension $(B, L, D)$ kontrahiert, bevor er zurück in das HBM geschrieben wird.
Zwischenzustände werden während des Backward-Passes im SRAM neu berechnet, anstatt im HBM gespeichert zu werden.

Ergebnisse

Synthetische Extrapolation: Bei der Induction-Heads-Aufgabe extrapoliert Mamba, trainiert auf Sequenzlänge 256, fehlerfrei (100 % Genauigkeit) bis zu einer Länge von $2^{20} = 1, 048, 576$ ( $4000 \times$ Sequenzlänge), während Attention-Baselines jenseits von $2^{14}$ versagen oder OOM (Out of Memory) laufen. Beim Selective Copying erreicht Mamba 99,8 % Genauigkeit im Vergleich zu 57,0 % bei H3 und 18,3 % bei S4.
Sprachmodellierung: Unter Chinchilla-Skalierung bis zu 1,3B Parametern auf The Pile evaluiert, erreicht oder übertrifft Mamba Transformer++ (Rotary Embeddings, SwiGLU, RMSNorm). Bei nachgelagerten Zero-Shot-Aufgaben (z. B. LAMBADA, PIQA, HellaSwag, WinoGrande, ARC) erzielt Mamba-2.8B eine durchschnittliche Genauigkeit von 63,3 % und übertrifft damit Pythia-2.8B (59,1 %) sowie Pythia-6.9B (61,7 %) und erreicht oder übertrifft Modelle mit 6B–7B Parametern.
Genomik und Audio: Bei der DNA-Klassifikation von Menschenaffen bei einer Sequenzlänge von 1M ( $2^{20}$ ) erreicht Mamba 71,67 % (1,4M Parameter) und 81,31 % (7M Parameter) Genauigkeit gegenüber 54,87 % bei HyenaDNA. Bei der SC09-Sprachgenerierung senkt Mamba die FID von 1,99 (SaShiMi) auf 0,94 (6,1M Parameter) und 0,67 (24,3M Parameter).
Effizienz: Der fusionierte Scan-Kernel läuft 20–40× schneller als standardmäßige PyTorch-Scan-Implementierungen und bis zu 7× schneller als FlashAttention-2 bei einem 32K-Kontext. Der Inferenzdurchsatz ist bis zu 5× höher als bei vergleichbaren Transformern.

Grenzen und offene Fragen

Leistung auf kontinuierlichen Signalen: Während eingabeabhängige Selektion das diskrete Schließen (Text, DNA) verbessert, zeigen Ablationen beim Vortraining auf rohen Audiowellenformen, dass rein zeitinvariante (LTI) Modelle (wie Standard-S4) auf gleichmäßig abgetasteten, kontinuierlichen Daten besser abschneiden.
Modellskalierung: Die evaluierten Modelle sind auf etwa 3B Parameter und 300B Pretraining-Tokens begrenzt; die empirische Leistung bei Skalierungen von 7B–70B+ und darüber hinaus erfordert weitere Untersuchungen.
Downstream-Ökosystem: Für Transformer existieren umfangreiche Werkzeuge für Prompt Engineering, In-Context-Manipulation, KV-Cache-Quantisierung und Adapter-Tuning. Wie sich SSMs an diese Interaktionsmodi anpassen lassen, bleibt Gegenstand aktiver Forschung.

Schlüsselbegriffe

Structured State Space Models (SSMs): Eine Klasse von tiefen Sequenzmodellen, die von kontinuierlichen linearen dynamischen Systemen abgeleitet sind und 1D-Eingaben über einen impliziten latenten Zustand abbilden.
Linear Time Invariance (LTI): Eine mathematische Eigenschaft, bei der die Systemübergangsparameter über alle Zeitschritte hinweg statisch bleiben, was schnelle globale Faltungsrepräsentationen ermöglicht, aber inhaltsabhängiges Schließen einschränkt.
Selection Mechanism: Die Parametrisierung der SSM-Übergangsdynamik (Δ, B, C) als Funktionen des Eingabe-Tokens, wodurch das Modell Kontext entlang der Sequenzlänge dynamisch weiterleiten, zurücksetzen oder ignorieren kann.
S6: Eine Abkürzung für selektive strukturierte Zustandsraummodelle, die über einen assoziativen parallelen Scan statt über Faltung berechnet werden.
Parallel Associative Scan: Ein arbeitseffizienter Präfix-Berechnungsalgorithmus, der lineare Rekurrenzen über Sequenzlängen hinweg parallel auf Beschleuniger-Hardware auswertet.
Hardware-Aware State Expansion: Eine Strategie aus GPU-Kernel-Fusion und Recomputation, die Diskretisierung und Scan im schnellen On-Chip-SRAM ausführt, um wiederholte Zugriffe auf große Zwischenzustände im High-Bandwidth-Memory (HBM) zu vermeiden.
Induction Heads: Ein synthetischer assoziativer Abrufmechanismus, der evaluiert, ob ein Modell ein zuvor beobachtetes Token nach dessen vorangehendem Kontext-Token reproduzieren kann.
Mamba Block: Ein vereinfachter, homogener Sequenzmodell-Block, der Eingabeprojektionen, 1D-Faltung, selektive SSMs und multiplikatives Gating ohne separate Attention- oder MLP-Schichten integriert.

Fragen

Warum können selektive SSMs nicht mittels schneller FFT-Faltungen trainiert werden?

Schnelle globale FFT-Faltungen erfordern, dass die Systemdynamik (die Matrizen A, B und C) über die Zeit konstant ist (lineare Zeitinvarianz). Da Mamba B, C und Δ an jedem Zeitschritt zu Funktionen des Eingabe-Tokens macht, variiert der Faltungskern an jeder Position. Dadurch bricht die Faltungsäquivalenz zusammen und erfordert stattdessen rekurrente parallele Scans.

Wie vermeidet Mamba die Speicherexplosion der rekurrenten Zustandsexpansion während des Trainings?

Naive Rekurrenz mit Zustandsexpansionsdimension N erfordert die Materialisierung eines Zwischenzustands der Form (B, L, D, N) im GPU-HBM. Mamba nutzt Kernel-Fusion, um nur niedrigdimensionale Eingaben in das SRAM zu laden, berechnet die Diskretisierung und den parallelen Scan im schnellen SRAM, schreibt nur die kontrahierte (B, L, D)-Ausgabe zurück in das HBM und berechnet Zwischenaktivierungen während der Backpropagation neu.

In welcher Beziehung steht Mambas Δ-Parameter zum rekurrenten Gating?

Der Schrittweitenparameter Δ fungiert als verallgemeinerter Gating-Mechanismus. Bei Diskretisierung unter Zero-Order-Hold (ZOH) mit N=1 und A=-1 reduziert sich die selektive Rekurrenz mathematisch auf ein Standard-RNN-Gate: Ein großes Δ veranlasst das Modell, sich auf die aktuelle Eingabe zu konzentrieren und den Zustand zurückzusetzen, während ein kleines Δ den historischen Zustand beibehält und die Eingabe ignoriert.

Verwendet Mamba während der Generierung einen KV-Cache?

Nein. Da Mamba während der Generierung als autoregressive Zustandsraum-Rekurrenz arbeitet, behält es lediglich einen verborgenen Zustand fester Größe (Dimension D × N pro Schicht) bei, anstatt Keys und Values für alle vergangenen Tokens zwischenzuspeichern, was zu einer Inferenz in konstanter Zeit O(1) pro Schritt führt.

Die Übersicht auf dieser Seite wurde von einer KI aus dem Paper erstellt und kann Fehler enthalten. Details stehen im Paper.

Mit Übersetzung und KI lesen