Mamba: Ridefinisce la Modellazione di Sequenze e Supera l’Architettura dei Transformer

In questo articolo su Mamba, esploreremo come questo innovativo modello nello spazio degli stati (SSM) rivoluziona 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 genomica e l’analisi audio. La sua modellazione delle sequenze in tempo lineare con spazi di stato selettivi garantisce prestazioni eccezionali in queste diverse modalità.

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

La vera unicità di Mamba risiede nella sua capacità di elaborazione rapida, nel suo strato 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 notevole nel machine learning.

Transformer vs Mamba

I Transformer, come GPT-4, hanno stabilito benchmark nel processamento 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 adatti a gestire sequenze di dati, come il testo per i modelli linguistici. A differenza dei modelli precedenti che elaboravano i dati in sequenza, 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 il prossimo elemento 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 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à autoregressiva.

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.

Ecco una panoramica di alto livello di come funziona un transformer:

1. Elaborazione input: I Transformer elaborano prima i dati di input in un formato che il modello possa comprendere, spesso utilizzando embedding che incorporano anche la posizione di ogni elemento nella sequenza.
2. Meccanismo di attenzione: Al suo nucleo, il meccanismo di attenzione calcola un punteggio che rappresenta quanto concentrarsi su altre parti della sequenza di input quando si comprende un elemento corrente.
3. Architettura encoder-decoder: Il modello Transformer è composto da un encoder per elaborare l’input e un decoder per generare l’output. Ciascuno consiste in più strati che raffinano la comprensione del modello dell’input.
4. Attenzione multi-testa: All’interno sia dell’encoder che del decoder, l’attenzione multi-testa consente al modello di concentrarsi contemporaneamente su diverse parti della sequenza da diversi spazi rappresentazionali, migliorando la sua capacità di apprendere da contesti diversi.
5. Rete neurale feed-forward posizionale: Dopo l’attenzione, una rete neurale semplice elabora l’output di ogni posizione separatamente e identicamente. Ciò viene combinato con l’input attraverso una connessione residua e seguito da normalizzazione del livello.
6. Generazione output: Il decoder predice quindi una sequenza di output, influenzata dal contesto dell’encoder e da ciò che ha generato finora.

La capacità del transformer di gestire sequenze in parallelo e il suo robusto meccanismo di attenzione lo rendono potente per compiti come la traduzione e la generazione di testo.

In contrasto, il modello Mamba opera in modo diverso utilizzando spazi di stato selettivi per elaborare le sequenze. Questo approccio affronta l’inefficienza computazionale dei Transformer quando si tratta di sequenze lunghe. Il design di Mamba consente un’inferenza più rapida e una scalabilità lineare con la lunghezza della sequenza, stabilendo un nuovo paradigma per la modellazione di sequenze che potrebbe essere più efficiente, specialmente man mano che le sequenze diventano sempre più lunghe.

Mamba

Mamba si distingue per l’abbandono dei blocchi tradizionali di attenzione e MLP. La struttura semplificata produce un modello più leggero e veloce, con un costo che cresce linearmente con la lunghezza della sequenza.

Le caratteristiche principali sono:

• SSM selettivi: filtrano le informazioni irrilevanti e concentrano l’elaborazione sui dati utili, favorendo il ragionamento efficiente basato sul contenuto.
• Algoritmo attento all’hardware: un algoritmo parallelo ottimizzato per l’hardware moderno, soprattutto le GPU, accelera il calcolo e riduce la memoria necessaria rispetto ai modelli tradizionali.
• Architettura semplificata: integra gli SSM selettivi ed elimina attenzione e blocchi MLP, offrendo una struttura più omogenea con migliore scalabilità e prestazioni.

Mamba ha ottenuto risultati elevati nel linguaggio, nell’audio e nella genomica, sia nel preaddestramento sia nei compiti specifici. Nella modellazione del linguaggio eguaglia o supera anche Transformer più grandi.

Codice e modelli preaddestrati sono disponibili pubblicamente su GitHub.

La copia standard è semplice per i modelli lineari. La copia selettiva e i compiti con induction head richiedono una memoria dinamica sensibile al contenuto.

I modelli strutturati nello spazio degli stati, o S4, sono una classe promettente che combina caratteristiche di RNN, CNN e modelli classici nello spazio degli stati. Si ispirano ai sistemi continui che trasformano funzioni o sequenze unidimensionali attraverso uno stato latente implicito. Nel deep learning offrono un metodo per costruire modelli sequenziali efficienti e adattabili.

La dinamica dei modelli S4

SSM (S4): il modello di base trasforma una sequenza x in un output y usando i parametri appresi A, B e C e il parametro temporale Δ. La trasformazione discretizza i parametri continui e applica l’operazione SSM. Nel caso invariante nel tempo, i parametri non cambiano fra i diversi passi temporali.

L’importanza della discretizzazione

La discretizzazione converte i parametri continui in parametri discreti mediante formule definite, mantenendo il collegamento con i sistemi a tempo continuo. Questo conferisce proprietà come l’invarianza alla risoluzione e una corretta normalizzazione, a beneficio di stabilità e prestazioni. Presenta inoltre analogie con i meccanismi di gating delle RNN, che regolano il flusso delle informazioni.

Linearità e invarianza temporale (LTI)

Una caratteristica fondamentale dei modelli S4 tradizionali è l’invarianza temporale lineare: le dinamiche restano costanti e i parametri sono fissi in tutti i passi. Questa proprietà consente ricorrenze e convoluzioni in un quadro semplice ma potente per la modellazione delle sequenze.

Superare i limiti fondamentali

La natura LTI limita la capacità di modellare dati che richiedono dinamiche adattive. Il lavoro su Mamba supera questo vincolo introducendo parametri variabili nel tempo, ampliando le sequenze e i compiti gestibili.

Il termine “modello nello spazio degli stati” comprende in generale processi ricorrenti con uno stato latente ed è usato in diverse discipline. Nel deep learning, S4 o SSM strutturati indica una classe specifica ottimizzata per il calcolo efficiente e la rappresentazione di sequenze complesse.

I modelli S4 possono essere integrati in reti neurali end-to-end come trasformazioni autonome di sequenze. Analogamente agli strati convoluzionali nelle CNN, possono costituire l’ossatura di diverse architetture sequenziali.

SSM e SSM con selezione a confronto

Perché introdurre la selettività

SSM strutturati

Secondo il lavoro, un aspetto fondamentale della modellazione sequenziale è comprimere il contesto in uno stato di dimensioni gestibili. Selezionare o filtrare gli input permette di conservare meglio tale stato, ottenendo modelli più efficienti ed efficaci. La selettività regola in modo adattivo il flusso di informazioni lungo la sequenza, una capacità importante per compiti linguistici complessi e altri ambiti.

Negli SSM selettivi i parametri dipendono dall’input, introducendo un’adattabilità assente nei modelli invarianti nel tempo. Ne risultano SSM variabili nel tempo che non possono più usare le convoluzioni nello stesso modo per il calcolo efficiente e ricorrono invece a un meccanismo ricorrente lineare.

SSM con selezione (S6): questa variante rende i parametri B, C e Δ dipendenti dall’input. Il modello può così concentrarsi selettivamente su parti della sequenza x. I parametri sono discretizzati tenendo conto della selezione e l’operazione SSM viene applicata in modo variabile nel tempo tramite una scansione che aggiorna dinamicamente lo stato lungo la sequenza.

Risultati principali di Mamba

Risultati di valutazione di Mamba riportati nell’articolo

Mamba combina velocità d’inferenza e accuratezza. Il progetto sfrutta meglio i contesti lunghi, come mostrano i risultati sulla modellazione del DNA e dell’audio, superando modelli precedenti in compiti con dipendenze a lungo raggio. La versatilità emerge anche nelle valutazioni senza esempi specifici su più compiti, con vantaggi di efficienza e scalabilità.

Iniziare a usare Mamba

I requisiti tecnici indicati all’epoca dell’articolo sono Linux, una GPU NVIDIA, PyTorch 1.12 o successivo e CUDA 11.6 o successivo. L’installazione usa comandi pip per i pacchetti del repository Mamba. In caso di problemi di compatibilità con PyTorch, il flag –no-build-isolation di pip può essere utile. I modelli, addestrati su dati come Pile e SlimPajama, rispondono a diverse esigenze computazionali e di valutazione.

Le interfacce vanno dallo strato SSM selettivo al blocco Mamba fino al modello linguistico completo. Il blocco Mamba è il modulo principale, usa uno strato Conv1d causale e può essere integrato in altre reti. L’articolo descrive come un esempio Python istanzi il modello ed elabori dati, sottolineando la semplicità e la flessibilità dell’interfaccia.

I modelli preaddestrati sono disponibili su Hugging Face, con dimensioni da 130 milioni a 2,8 miliardi di parametri e addestramento su Pile e SlimPajama. Sono progettati per diverse esigenze di calcolo e prestazioni, rispettando gli standard dimensionali di GPT-3. L’elevato throughput e l’accuratezza ne fanno un’opzione competitiva per la modellazione del linguaggio e altre applicazioni.

L’impatto di Mamba

Mamba offre un’alternativa ai Transformer per elaborare dati ricchi di informazioni. Il progetto sfrutta l’hardware moderno ottimizzando memoria ed elaborazione parallela. La disponibilità open source del codice e dei modelli preaddestrati lo rende accessibile a ricercatori e sviluppatori di AI e deep learning.