Mamba: Lineaire-tijdsequentiemodellering met selectieve toestandsruimten

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

December 1, 2023

Auteurs: Albert Gu, Tri Dao

cs.AI

Samenvatting

Foundation models, die momenteel de meeste spannende toepassingen in deep learning aandrijven, zijn bijna universeel gebaseerd op de Transformer-architectuur en de kernmodule ervan, attention. Veel subkwadratische-tijdarchitecturen zoals lineaire attention, gated convolution en recurrente modellen, en gestructureerde state space-modellen (SSM's) zijn ontwikkeld om de computationele inefficiëntie van Transformers op lange sequenties aan te pakken, maar ze presteren niet zo goed als attention op belangrijke modaliteiten zoals taal. Wij identificeren dat een belangrijk zwaktepunt van dergelijke modellen hun onvermogen is om op inhoud gebaseerd redeneren uit te voeren, en introduceren verschillende verbeteringen. Ten eerste wordt door simpelweg de SSM-parameters functies van de invoer te laten zijn, hun zwakte met discrete modaliteiten aangepakt, waardoor het model selectief informatie kan doorgeven of vergeten langs de sequentielengtedimensie, afhankelijk van het huidige token. Ten tweede, hoewel deze verandering het gebruik van efficiënte convoluties verhindert, ontwerpen we een hardwarebewust parallel algoritme in recurrente modus. We integreren deze selectieve SSM's in een vereenvoudigd end-to-end neuraal netwerkarchitectuur zonder attention of zelfs MLP-blokken (Mamba). Mamba geniet van snelle inferentie (5 keer hogere doorvoer dan Transformers) en lineaire schaalbaarheid in sequentielengte, en de prestaties verbeteren op echte data tot sequenties van miljoen lengtes. Als een algemene sequentiemodel-backbone bereikt Mamba state-of-the-art prestaties over verschillende modaliteiten zoals taal, audio en genomica. Op taalmodellering presteert ons Mamba-3B-model beter dan Transformers van dezelfde grootte en evenaart het Transformers die twee keer zo groot zijn, zowel in pretraining als in downstream-evaluatie.

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.