Mamba: Lineare Zeit-Sequence-Modellierung mit selektiven Zustandsräumen

Grundmodellarchitekturen, die derzeit die meisten spannenden Anwendungen im Bereich des tiefen Lernens antreiben, basieren nahezu universell auf der Transformer-Architektur und ihrem zentralen Aufmerksamkeitsmodul. Zahlreiche Architekturen mit subquadratischer Laufzeit – wie lineare Aufmerksamkeit, gated Convolutionen, rekurrente Modelle sowie strukturierte Zustandsraummodelle (SSMs) – wurden entwickelt, um die rechnerische Ineffizienz von Transformers bei langen Sequenzen zu überwinden. Doch sie erzielen auf wichtigen Modalitäten wie Sprache nicht denselben Erfolg wie die klassische Aufmerksamkeit. Wir identifizieren, dass ein zentraler Schwachpunkt solcher Modelle ihre Unfähigkeit ist, inhaltsbasierte Schlussfolgerungen zu ziehen, und stellen mehrere Verbesserungen vor. Erstens ermöglicht die Erweiterung der SSM-Parameter auf Funktionen der Eingabeadressen eine bessere Bewältigung diskreter Modalitäten, wodurch das Modell gezielt Informationen entlang der Sequenzlänge propagieren oder vergessen kann, abhängig vom aktuellen Token. Zweitens, obwohl diese Änderung die Nutzung effizienter Faltungen verhindert, entwerfen wir einen hardwarebewussten parallelen Algorithmus im rekurrenten Modus. Wir integrieren diese selektiven SSMs in eine vereinfachte, end-to-end neuronale Architektur ohne Aufmerksamkeits- oder gar MLP-Blöcke (Mamba). Mamba zeichnet sich durch schnelle Inferenz aus (5-mal höhere Durchsatzrate als Transformers) und zeigt eine lineare Skalierbarkeit mit der Sequenzlänge. Seine Leistung steigert sich bei realen Daten bis hin zu Sequenzen mit einer Länge von einer Million. Als allgemeiner Sequenzmodell-Backbone erreicht Mamba state-of-the-art Ergebnisse in mehreren Modalitäten wie Sprache, Audio und Genomik. Bei der Sprachmodellierung übertrifft unser Mamba-3B-Modell Transformers gleicher Größe und erreicht die Leistung von Transformers mit doppelter Größe – sowohl im Pretraining als auch in nachfolgenden Evaluierungen.