Mamba Ridefinendo la modellazione sequenziale e superando l’architettura dei trasformatori

Mamba Ripensando la modellazione sequenziale e superando l'architettura dei trasformatori

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

Approfondiremo la capacità di Mamba di superare le sfide computazionali affrontate dai tradizionali Transformer, soprattutto con sequenze lunghe. Il suo approccio selettivo nei modelli di spazio di stato consente un’infusione più veloce e una scalabilità lineare rispetto alla lunghezza della sequenza, migliorando significativamente la velocità di elaborazione.

L’unicità di Mamba risiede nella sua capacità di elaborazione rapida, nello strato selettivo di SSM e nel design ottimizzato per l’hardware, ispirato a FlashAttention. Queste caratteristiche consentono a Mamba di superare molti modelli esistenti, compresi quelli basati sull’approccio dei Transformer, rendendolo un notevole avanzamento nell’apprendimento automatico.

Transformer vs Mamba 

I Transformer, come GPT-4, hanno stabilito nuovi standard nell’elaborazione del linguaggio naturale. Tuttavia, la loro efficienza diminuisce con sequenze più lunghe. Qui entra in gioco Mamba, con la sua capacità di elaborare sequenze lunghe in modo più efficiente e la sua architettura unica che semplifica l’intero processo.

• Annunciamo i Summit sull’Ingegneria dei Dati e l’Intelligenza Artificiale per l’Innovazione del 2024′.
• Trey Doig, CTO e Co-Fondatore di Pathlight – Serie di interviste
• Recensione di Sudowrite L’IA può scrivere un romanzo che suona umano?

I Transformer sono abili nell’affrontare sequenze di dati, come il testo per i modelli di linguaggio. A differenza dei modelli precedenti che elaboravano i dati in modo sequenziale, i Transformer elaborano intere sequenze contemporaneamente, consentendo di catturare relazioni complesse tra i dati.

Utilizzano un meccanismo di attenzione che consente al modello di concentrarsi su diverse parti della sequenza durante la previsione.

Questa attenzione viene calcolata utilizzando tre insiemi di pesi: query, chiavi e valori, derivati dai dati di input. Ciascun elemento di una sequenza viene confrontato con ogni altro elemento, fornendo un peso che indica l’importanza, o “attenzione”, che ogni elemento dovrebbe ricevere nella previsione del prossimo elemento nella sequenza.

I Transformer comprendono due blocchi principali: l’encoder, che elabora i dati di input, e il decoder, che genera l’output. L’encoder è composto da più strati, ciascuno contenente due sub-layers: meccanismo di auto-attenzione multi-head e una semplice rete di feed-forward connesse in modo posizionale. Normalizzazione e connessioni residue vengono utilizzate in ciascun sub-layer per facilitare l’addestramento delle reti profonde.

Il decoder ha anche strati con due sub-layers simili all’encoder, ma aggiunge un terzo sub-layer che effettua l’attenzione multi-head sull’output dell’encoder. La natura sequenziale del decoder garantisce che le previsioni per una posizione possano considerare solo le posizioni precedenti, preservando la proprietà autoregressiva.

Contrariamente ai 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, offrendo un’elaborazione più efficiente per le sequenze. Il design di Mamba consente un’infusione più veloce e una scalabilità lineare rispetto alla lunghezza della sequenza, aprendo una nuova era nella modellazione delle sequenze che potrebbe risultare più efficiente, soprattutto assequenze sempre più lunghe.

Ecco una panoramica generale di come funziona un transformer:

1. Elaborazione dell’input: I transformer codificano prima i dati di input in un formato comprensibile per il modello, spesso utilizzando embedding che incorporano anche la posizione di ciascun elemento nella sequenza.
2. Meccanismo di attenzione: Al suo nucleo, il meccanismo di attenzione calcola uno score che rappresenta quanto focalizzare altre parti della sequenza di input durante la comprensione di 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. Ognuno di essi è composto da più strati che affinano la comprensione del modello dell’input.
4. Attenzione multi-head: Sia nell’encoder che nel decoder, l’attenzione multi-head consente al modello di concentrarsi contemporaneamente su diverse parti della sequenza da rappresentazioni spaziali diverse, migliorando la sua capacità di apprendere da contesti diversi.
5. Reti di feed-forward posizionali: Dopo l’attenzione, una semplice rete neurale elabora l’output di ogni posizione separatamente e in modo identico. Ciò viene combinato con l’input attraverso una connessione residua e seguito da una normalizzazione di livello.
6. Generazione dell’output: Il decoder quindi predice 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 funziona in modo diverso utilizzando spazi di stato selettivi per elaborare sequenze. Questo approccio affronta l’inefficienza computazionale dei Transformer quando si tratta di sequenze lunghe. Il design di Mamba consente un’infusione più veloce e una scalabilità lineare rispetto alla lunghezza della sequenza, aprendo un nuovo paradigma per la modellazione delle sequenze che potrebbe essere più efficiente, specialmente con l’aumentare della lunghezza delle sequenze.

Mamba

Ciò che rende Mamba veramente unico è la sua differenza rispetto ai tradizionali blocchi di attenzione e MLP. Questa semplificazione porta ad un modello più leggero e più veloce che scala linearmente con la lunghezza della sequenza, un risultato senza pari rispetto ai suoi predecessori.

Le caratteristiche principali di Mamba includono:

1. SSM selettive: Queste permettono a Mamba di filtrare informazioni irrilevanti e concentrarsi su dati pertinenti, migliorando il trattamento delle sequenze. Questa selettività è cruciale per un ragionamento basato sul contenuto efficiente.
2. Algoritmo consapevole dell’hardware: Mamba utilizza un algoritmo parallelo ottimizzato per hardware moderno, in particolare per le GPU. Questo design consente un calcolo più rapido e riduce i requisiti di memoria rispetto ai modelli tradizionali.
3. Architettura semplificata: Integrando SSM selettive ed eliminando blocchi di attenzione e MLP, Mamba offre una struttura più semplice e omogenea. Ciò porta a una migliore scalabilità e prestazioni.

Mamba ha dimostrato prestazioni superiori in vari settori, compresi il linguaggio, l’audio e la genomica, eccellendo sia nel pretraining che nelle attività specifiche di dominio. Ad esempio, nel modellamento del linguaggio, Mamba eguaglia o supera le prestazioni di modelli Transformer più grandi.

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

I compiti di copia standard sono semplici per i modelli lineari. La copia selettiva e le Teste di Induzione richiedono memoria dinamica e consapevole del contenuto per i LLM.

I modelli di Spazio di Stato Strutturato (S4) sono emersi di recente come una classe promettente di modelli di sequenza, che comprendono caratteristiche di RNN, CNN e modelli classici di spazio di stato. I modelli S4 traggono ispirazione da sistemi continui, in particolare un tipo di sistema che mappa funzioni o sequenze unidimensionali attraverso uno stato latente implicito. Nel contesto dell’apprendimento approfondito, rappresentano un’innovazione significativa, fornendo una nuova metodologia per la progettazione di modelli di sequenza efficienti e altamente 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’uscita y utilizzando i parametri appresi A, B, C e un parametro di ritardo Δ. La trasformazione coinvolge la discretizzazione dei parametri (trasformando funzioni continue in funzioni discrete) e l’applicazione dell’operazione SSM, che è time-invariante, il che significa che non cambia con i diversi passaggi di tempo.

La rilevanza della discretizzazione

La discretizzazione è un processo chiave che trasforma i parametri continui in parametri discreti attraverso formule fisse, consentendo ai modelli S4 di mantenere una connessione con i sistemi a tempo continuo. Ciò conferisce ai modelli proprietà aggiuntive, come l’invarianza della risoluzione, e garantisce una corretta normalizzazione, migliorando la stabilità e le prestazioni del modello. La discretizzazione richiama anche i meccanismi di gating trovati nelle RNN, che sono fondamentali per gestire il flusso delle informazioni attraverso la rete.

L’Invarianza Lineare nel Tempo (LTI)

Una caratteristica fondamentale dei modelli S4 è la loro invarianza lineare nel tempo. Questa proprietà implica che la dinamica del modello rimane consistente nel tempo, con i parametri fissi per tutti gli intervalli di tempo. LTI è una base fondamentale per la ricorrenza e le convoluzioni, offrendo un framework semplificato ma potente per la costruzione di modelli di sequenza.

Superare le Limitazioni Fondamentali

Il framework S4 è tradizionalmente limitato dalla sua natura LTI, che rappresenta una sfida nella modellazione di dati che richiedono dinamiche adattive. Il recente articolo di ricerca presenta un approccio che supera queste limitazioni introducendo parametri variabili nel tempo, rimuovendo così il vincolo dell’LTI. Ciò consente ai modelli S4 di gestire un insieme più ampio di sequenze e compiti, ampliandone significativamente l’applicabilità.

Il termine “modello di spazio di stato” copre in generale ogni processo ricorrente che coinvolge uno stato latente ed è stato utilizzato per descrivere vari concetti in diversi settori. Nel contesto dell’apprendimento approfondito, i modelli S4, o SSM strutturati, si riferiscono a una classe specifica di modelli che sono stati ottimizzati per un calcolo efficiente pur mantenendo la capacità di modellare sequenze complesse.

Gli modelli S4 possono essere integrati nelle architetture di reti neurali end-to-end, funzionando come trasformazioni di sequenze indipendenti. Possono essere visti come analoghi ai livelli di convoluzione nelle CNN, fornendo la struttura portante per la modellazione delle sequenze in una varietà di architetture di reti neurali.

SSM vs SSM + Selezione

Motivazione per la Selettività nella Modellazione delle Sequenze

SSM Strutturato

Il documento sostiene che un aspetto fondamentale della modellazione delle sequenze è la compressione del contesto in uno stato gestibile. I modelli che possono selezionare o filtrare in modo selettivo gli input forniscono un modo più efficace di mantenere questo stato compresso, portando a modelli di sequenze più efficienti e potenti. Questa selettività è fondamentale affinché i modelli possano controllare in modo adattivo come l’informazione fluisce lungo la dimensione della sequenza, una capacità essenziale per gestire compiti complessi nella modellazione del linguaggio e oltre.

Gli SSM selettivi potenziano gli SSM convenzionali consentendo che i loro parametri siano dipendenti dagli input, introducendo così un grado di adattabilità precedentemente irraggiungibile con modelli tempo-invarianti. Ciò si traduce in SSM che variano nel tempo e che non possono più utilizzare le convoluzioni per un calcolo efficiente, ma si basano invece su un meccanismo di ricorrenza lineare, una deviazione significativa dai modelli tradizionali.

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

Prestazioni di Rilievo di Mamba

Mamba è il migliore in ogni singolo risultato di valutazione

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 nella modellazione del DNA e dell’audio, superando i modelli precedenti in compiti complessi che richiedono dipendenze a lungo raggio. La sua versatilità è evidenziata anche nelle valutazioni zero-shot su vari compiti, ponendo un nuovo standard per questo tipo di 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 prevede semplici comandi pip per installare i pacchetti necessari dal repository di Mamba. Se si verificano problemi di compatibilità con le versioni di PyTorch, l’utilizzo del flag –no-build-isolation con pip può essere d’aiuto. Questi modelli, addestrati su ampi set di dati come il Pile e il dataset SlimPajama, sono progettati per soddisfare diverse esigenze computazionali e benchmark di prestazioni.

Mamba offre diversi livelli di interfacce, dal livello di selezione SSM al blocco di Mamba e alle strutture complete del modello di linguaggio. Il blocco di Mamba, che rappresenta il modulo principale dell’architettura, utilizza un livello Conv1d causale e può essere facilmente integrato in design di reti neurali. L’esempio di utilizzo fornito in Python illustra come istanziare un modello di Mamba e elaborare i dati attraverso di esso, mettendo in evidenza la semplicità e la flessibilità del sistema.

I modelli pre-addestrati di Mamba sono disponibili su Hugging Face/a>, con dimensioni che vanno da 130M a 2.8B di parametri, addestrati sugli ampi dataset Pile e SlimPajama. Questi modelli sono progettati per soddisfare diverse esigenze computazionali e di prestazioni, rispettando gli standard dimensionali di GPT-3. Gli utenti possono aspettarsi un’elevata velocità di elaborazione e accuratezza da questi modelli, rendendo Mamba una scelta competitiva per diverse applicazioni, tra cui ma non limitate alla modellazione del linguaggio.

Impatto di Mamba

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