Mamba: Ridefinizione della modellazione di sequenze e superamento dell’architettura dei Transformer

In questo articolo su Mamba, esploreremo come questo innovativo modello di spazio di stato (SSM) rivoluzioni la modellazione di sequenze. Sviluppato da Albert Gu e Tri Dao, Mamba si distingue per la sua efficienza nel processare sequenze complesse in campi come l’elaborazione del linguaggio, la genomics e l’analisi audio. La sua modellazione di sequenze in tempo lineare con spazi di stato selettivi garantisce prestazioni eccezionali in queste diverse modalità.

Esamineremo la capacità di Mamba di superare le sfide computazionali affrontate dai tradizionali Transformer, in particolare con sequenze lunghe. Il suo approccio selettivo nei modelli di spazio di stato consente un’inferenza più rapida e una scalabilità lineare con la lunghezza della sequenza, migliorando notevolmente il throughput.

L’unicità di Mamba risiede nella sua capacità di elaborazione rapida, nel layer SSM selettivo e nel design hardware-friendly ispirato a FlashAttention. Queste caratteristiche consentono a Mamba di superare molti modelli esistenti, compresi quelli basati sull’approccio Transformer, rendendolo un progresso significativo nel campo dell’apprendimento automatico.

Transformer vs Mamba

I Transformer, come GPT-4, hanno stabilito benchmark nell’elaborazione del linguaggio naturale. Tuttavia, la loro efficienza cala con sequenze più lunghe. È qui che Mamba si distingue, con la sua capacità di elaborare sequenze lunghe in modo più efficiente e la sua architettura unica che semplifica l’intero processo.

I Transformer sono abili nel gestire sequenze di dati, come il testo per i modelli linguistici. A differenza dei modelli precedenti che elaboravano i dati in modo sequenziale, i Transformer elaborano l’intera sequenza contemporaneamente, consentendo loro di catturare relazioni complesse all’interno dei dati.

Utilizzano un meccanismo di attenzione, che consente al modello di concentrarsi su diverse parti della sequenza quando si effettuano previsioni.

Questa attenzione viene calcolata utilizzando tre set di pesi: query, chiavi e valori, derivati dai dati di input. Ogni elemento in una sequenza viene confrontato con ogni altro elemento, fornendo un peso che segnala l’importanza, o ‘attenzione’, che ogni elemento dovrebbe ricevere quando si prevede l’elemento successivo nella sequenza.

I Transformer mantengono due blocchi principali: l’encoder, che elabora i dati di input, e il decoder, che genera l’output. L’encoder consiste in più strati, ciascuno contenente due sottolayer: un meccanismo di auto-attenzione multi-testa e una rete neurale feed-forward fully connected semplice e posizionale. La normalizzazione e le connessioni residue vengono utilizzate in ogni sottolayer per aiutare nell’addestramento di reti profonde.

Il decoder ha anche strati con due sottolayer simili all’encoder, ma aggiunge un terzo sottolayer che esegue un’attenzione multi-testa sull’output dell’encoder. La natura sequenziale del decoder garantisce che le previsioni per una posizione possano considerare solo posizioni precedenti, preservando la proprietà auto-regressiva.

In contrasto con i Transformer, il modello Mamba adotta un approccio diverso. Mentre i Transformer affrontano il problema delle sequenze lunghe utilizzando meccanismi di attenzione più complessi, Mamba utilizza spazi di stato selettivi, fornendo un approccio più computazionalmente efficiente.

Mamba si distingue per la sua capacità di elaborazione rapida, il layer SSM selettivo e il design hardware-friendly ispirato a FlashAttention. Queste caratteristiche consentono a Mamba di superare molti modelli esistenti, compresi quelli basati sull’approccio Transformer, rendendolo un progresso significativo nel campo dell’apprendimento automatico.

Mamba

Ciò che rende Mamba veramente unico è la sua deviazione dai tradizionali blocchi di attenzione e MLP. Questa semplificazione porta a un modello più leggero e veloce che scala linearmente con la lunghezza della sequenza – una caratteristica senza precedenti rispetto ai suoi predecessori.

Le caratteristiche chiave di Mamba includono:

1. SSM selettivi: questi consentono a Mamba di filtrare le informazioni non rilevanti e concentrarsi sui dati rilevanti, migliorando la sua gestione delle sequenze. Questa selettività è cruciale per il ragionamento basato sul contenuto efficiente.
2. Algoritmo hardware-aware: Mamba utilizza un algoritmo parallelo ottimizzato per l’hardware moderno, in particolare le GPU. Questo design consente calcoli più rapidi e riduce i requisiti di memoria rispetto ai modelli tradizionali.
3. Architettura semplificata: integrando SSM selettivi ed eliminando i blocchi di attenzione e MLP, Mamba offre una struttura più semplice e omogenea. Ciò porta a una migliore scalabilità e prestazione.

Mamba ha dimostrato prestazioni superiori in vari domini, tra cui linguaggio, audio e genomics, eccellendo sia nella pre-formazione che nei compiti specifici del dominio.

Il codice e i modelli pre-addestrati di Mamba sono disponibili apertamente per l’uso della comunità su GitHub.

Standard Copying tasks are simple for linear models. Selective Copying and Induction Heads require dynamic, content-aware memory for LLMs.

I modelli di spazio di stato strutturati (S4) sono recentemente emersi come una classe promettente di modelli di sequenze, che combinano caratteristiche di RNN, CNN e modelli di spazio di stato classici. I modelli S4 traggono ispirazione da sistemi continui, in particolare da un tipo di sistema che mappa funzioni o sequenze unidimensionali attraverso uno stato latente implicito. Nel contesto dell’apprendimento profondo, rappresentano un’innovazione significativa, fornendo una nuova metodologia per la progettazione di modelli di sequenze efficienti e adattabili.

La dinamica dei modelli S4

SSM (S4) questo è il modello di spazio di stato strutturato di base. Prende una sequenza x e produce un output y utilizzando parametri appresi A, B, C e un parametro di ritardo Δ. La trasformazione coinvolge la discretizzazione dei parametri (trasformando funzioni continue in discrete) e l’applicazione dell’operazione SSM, che è invariante nel tempo – ovvero non cambia nel corso dei diversi passaggi temporali.

L’importanza della discretizzazione

La discretizzazione è un processo fondamentale che trasforma i parametri continui in discreti attraverso formule fisse, consentendo ai modelli S4 di mantenere una connessione con i sistemi continui nel tempo. Ciò conferisce ai modelli proprietà aggiuntive, come l’invarianza di risoluzione, e garantisce una corretta normalizzazione, migliorando la stabilità e le prestazioni del modello. La discretizzazione traccia anche paralleli con i meccanismi di gating presenti nelle RNN, che sono cruciali per la gestione del flusso di informazioni attraverso la rete.

Linearità nel tempo (LTI)

Una caratteristica fondamentale dei modelli S4 è la loro linearità nel tempo. Ciò implica che la dinamica del modello rimane coerente nel tempo, con i parametri fissi per tutti i passaggi temporali. LTI è un pilastro della ricorrenza e delle convoluzioni, offrendo un framework semplificato ma potente per la costruzione di modelli di sequenze.

Superare le limitazioni fondamentali

La cornice S4 è stata tradizionalmente limitata dalla sua natura LTI, che pone sfide nella modellazione di dati che richiedono dinamiche adattive. La recente ricerca presenta un approccio che supera queste limitazioni introducendo parametri che variano nel tempo, rimuovendo così la limitazione LTI. Ciò consente ai modelli S4 di gestire un insieme più diversificato di sequenze e compiti, ampliando notevolmente la loro applicabilità.

Il termine ‘modello di spazio di stato’ copre in generale qualsiasi processo ricorrente che coinvolga uno stato latente e è stato utilizzato per descrivere vari concetti in diverse discipline. Nel contesto dell’apprendimento profondo, i modelli S4, o SSM strutturati, si riferiscono a una classe specifica di modelli che sono stati ottimizzati per il calcolo efficiente pur mantenendo la capacità di modellare sequenze complesse.

I modelli S4 possono essere integrati in architetture di reti neurali end-to-end, funzionando come trasformazioni di sequenze autonome. Possono essere visti come analoghi ai layer di convoluzione nelle CNN, fornendo la colonna vertebrale per la modellazione di sequenze in una varietà di architetture di reti neurali.

SSM vs SSM + Selection

Motivazione per la selettività nella modellazione di sequenze

Structured SSMs

Il paper sostiene che un aspetto fondamentale della modellazione di sequenze è la compressione del contesto in uno stato gestibile. I modelli che possono selettivamente concentrarsi o filtrare gli input forniscono un mezzo più efficace per mantenere questo stato compresso, portando a modelli di sequenze più efficienti e potenti. Questa selettività è vitale per i modelli che adattivamente controllano come le informazioni fluiscono lungo la dimensione della sequenza, una capacità essenziale per gestire compiti complessi nella modellazione del linguaggio e oltre.

Gli SSM selettivi migliorano gli SSM convenzionali consentendo ai loro parametri di essere dipendenti dall’input, introducendo un grado di adattività precedentemente irraggiungibile con modelli a tempo invariante. Ciò porta a SSM a tempo variabile che non possono più utilizzare le convoluzioni per il calcolo efficiente, ma si affidano invece a un meccanismo di ricorrenza lineare, una deviazione significativa dai modelli tradizionali.

SSM + Selezione (S6) questa variante include un meccanismo di selezione, aggiungendo una dipendenza dall’input ai parametri B e C e un parametro di ritardo Δ. Ciò consente al modello di concentrarsi selettivamente su parti specifiche della sequenza di input x. I parametri vengono discretizzati tenendo conto della selezione e l’operazione SSM viene applicata in modo tempo-variabile utilizzando un’operazione di scansione, che elabora gli elementi in modo sequenziale, regolando la concentrazione dinamicamente nel corso del tempo.

Prestazioni di Mamba

Mamba is best-in-class on every single evaluation result

In termini di prestazioni, Mamba eccelle sia nella velocità di inferenza che nell’accuratezza. Il suo design consente un miglior utilizzo di contesti più lunghi, come dimostrato sia nella modellazione del DNA che in quella dell’audio, superando i modelli precedenti in compiti complessi che richiedono dipendenze a lungo raggio. La sua versatilità è anche evidenziata nelle valutazioni zero-shot across multiple task, stabilendo un nuovo standard per tali modelli in termini di efficienza e scalabilità.

Iniziare con Mamba

Per coloro che sono interessati a sfruttare Mamba, i requisiti tecnici includono un sistema operativo Linux, una GPU NVIDIA, PyTorch 1.12+ e CUDA 11.6+. L’installazione comporta semplici comandi pip per installare i pacchetti necessari dal repository Mamba. Se si verificano problemi di compatibilità con le versioni di PyTorch, utilizzare il flag –no-build-isolation con pip può aiutare. Questi modelli, addestrati su dataset estensivi come Pile e SlimPajama, sono progettati per soddisfare diverse esigenze computazionali e benchmark di prestazione.

Mamba offre diversi livelli di interfacce, dal layer SSM selettivo al blocco Mamba e alle strutture complete del modello linguistico. Il blocco Mamba, che è il modulo principale dell’architettura, utilizza un layer Conv1d causale e può essere facilmente integrato nei progetti di reti neurali. L’esempio di utilizzo in Python dimostra l’istanziazione di un modello Mamba e l’elaborazione dei dati attraverso di esso, evidenziando la semplicità e la flessibilità del sistema.

I modelli pre-addestrati di Mamba sono disponibili su Hugging Face, con dimensioni che vanno da 130M a 2,8B parametri, addestrati sul dataset Pile estensivo e sul dataset SlimPajama. Questi modelli sono progettati per soddisfare diverse esigenze computazionali e di prestazione, aderendo agli standard dimensionali di GPT-3. Gli utenti possono aspettarsi un alto throughput e accuratezza da questi modelli, rendendo Mamba una scelta competitiva per varie applicazioni, tra cui, ma non limitato a, la modellazione del linguaggio.

Impatto di Mamba

Mamba rappresenta un balzo in avanti nella modellazione di sequenze, offrendo un’alternativa potente alle architetture Transformer per l’elaborazione di dati densi di informazioni. Il suo design si allinea con le esigenze dell’hardware moderno, ottimizzando sia l’utilizzo della memoria che le capacità di elaborazione parallela. La disponibilità open-source del codice e dei modelli pre-addestrati di Mamba lo rende uno strumento accessibile e robusto per ricercatori e sviluppatori nel campo dell’IA e dell’apprendimento profondo.