Mamba: Modellazione di Sequenze in Tempo Lineare con Spazi di Stato Selettivi

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

December 1, 2023

Autori: Albert Gu, Tri Dao

cs.AI

Abstract

I modelli di base, che oggi alimentano la maggior parte delle applicazioni più interessanti nel deep learning, sono quasi universalmente basati sull'architettura Transformer e sul suo modulo di attenzione centrale. Sono state sviluppate molte architetture a tempo subquadratico, come l'attenzione lineare, i modelli convoluzionali con gate e ricorrenti, e i modelli strutturati a spazio di stati (SSM), per affrontare l'inefficienza computazionale dei Transformer su sequenze lunghe, ma non hanno performato altrettanto bene quanto l'attenzione su modalità importanti come il linguaggio. Identifichiamo che una debolezza chiave di tali modelli è la loro incapacità di eseguire ragionamenti basati sul contenuto, e apportiamo diversi miglioramenti. Innanzitutto, semplicemente rendendo i parametri dello SSM funzioni dell'input si affronta la loro debolezza con modalità discrete, permettendo al modello di propagare o dimenticare selettivamente informazioni lungo la dimensione della lunghezza della sequenza in base al token corrente. In secondo luogo, anche se questo cambiamento impedisce l'uso di convoluzioni efficienti, progettiamo un algoritmo parallelo consapevole dell'hardware in modalità ricorrente. Integriamo questi SSM selettivi in un'architettura di rete neurale end-to-end semplificata senza attenzione o persino blocchi MLP (Mamba). Mamba gode di un'inferenza veloce (5 volte più veloce dei Transformer) e di una scalabilità lineare nella lunghezza della sequenza, e le sue prestazioni migliorano su dati reali fino a sequenze di lunghezza milioni. Come modello di base per sequenze generali, Mamba raggiunge prestazioni all'avanguardia su diverse modalità come linguaggio, audio e genomica. Nel modeling del linguaggio, il nostro modello Mamba-3B supera i Transformer della stessa dimensione e eguaglia i Transformer di dimensioni doppie, sia nel pre-training che nella valutazione downstream.

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.