I Transformer sono SSM: Modelli Generalizzati e Algoritmi Efficienti Attraverso la Dualità degli Spazi di Stato Strutturati

Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality

May 31, 2024

Autori: Tri Dao, Albert Gu

cs.AI

Abstract

Mentre i Transformer sono stati l'architettura principale dietro il successo del deep learning nel campo della modellazione del linguaggio, i modelli a spazio di stati (SSM) come Mamba hanno recentemente dimostrato di eguagliare o superare i Transformer su scala piccola e media. Mostriamo che queste famiglie di modelli sono in realtà strettamente correlate e sviluppiamo un ricco quadro di connessioni teoriche tra gli SSM e le varianti dell'attenzione, collegate attraverso varie decomposizioni di una classe ben studiata di matrici semiseparabili strutturate. Il nostro framework di dualità dello spazio di stati (SSD) ci permette di progettare una nuova architettura (Mamba-2) il cui strato centrale è un perfezionamento dell'SSM selettivo di Mamba, che risulta da 2 a 8 volte più veloce, continuando a essere competitivo con i Transformer nella modellazione del linguaggio.

English

While Transformers have been the main architecture behind deep learning's success in language modeling, state-space models (SSMs) such as Mamba have recently been shown to match or outperform Transformers at small to medium scale. We show that these families of models are actually quite closely related, and develop a rich framework of theoretical connections between SSMs and variants of attention, connected through various decompositions of a well-studied class of structured semiseparable matrices. Our state space duality (SSD) framework allows us to design a new architecture (Mamba-2) whose core layer is an a refinement of Mamba's selective SSM that is 2-8X faster, while continuing to be competitive with Transformers on language modeling.