Mamba: Lineare Sequenzmodellierung mit selektiven Zustandsräumen

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

December 1, 2023

Autoren: Albert Gu, Tri Dao

cs.AI

Zusammenfassung

Foundation Models, die heute die meisten der spannenden Anwendungen im Bereich des Deep Learning antreiben, basieren fast ausschließlich auf der Transformer-Architektur und ihrem zentralen Attention-Modul. Viele subquadratische Architekturen wie lineare Attention, gated Convolution und rekurrente Modelle sowie strukturierte State-Space-Modelle (SSMs) wurden entwickelt, um die rechenineffiziente Verarbeitung langer Sequenzen durch Transformers zu adressieren, aber sie konnten nicht die gleiche Leistung wie Attention in wichtigen Modalitäten wie Sprache erzielen. Wir identifizieren, dass eine zentrale Schwäche solcher Modelle ihre Unfähigkeit zur inhaltsbasierten Schlussfolgerung ist, und schlagen mehrere Verbesserungen vor. Erstens adressiert die einfache Anpassung der SSM-Parameter als Funktionen des Inputs ihre Schwäche bei diskreten Modalitäten, wodurch das Modell selektiv Informationen entlang der Sequenzlängendimension propagieren oder vergessen kann, abhängig vom aktuellen Token. Zweitens, obwohl diese Änderung die Verwendung effizienter Faltungen verhindert, entwickeln wir einen hardwarebewussten parallelen Algorithmus im rekurrenten Modus. Wir integrieren diese selektiven SSMs in eine vereinfachte end-to-end neuronale Netzarchitektur ohne Attention oder sogar MLP-Blöcke (Mamba). Mamba ermöglicht schnelle Inferenz (5-mal höherer Durchsatz als Transformers) und lineare Skalierung der Sequenzlänge, und seine Leistung verbessert sich bei realen Daten bis zu Sequenzen mit Millionen von Elementen. Als allgemeines Sequenzmodell erreicht Mamba state-of-the-art Leistung in mehreren Modalitäten wie Sprache, Audio und Genomik. Beim Sprachmodellieren übertrifft unser Mamba-3B-Modell Transformers gleicher Größe und erreicht die Leistung von Transformers, die doppelt so groß sind, sowohl im Pretraining als auch in der Downstream-Evaluation.

English

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 (5times 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.