Intelligenza Artificiale
Mamba: ridefinire la modellazione delle sequenze e migliorare l'architettura dei trasformatori
In questo articolo su Mamba, esploreremo come questo innovativo modello di stato-spazio (SSM) rivoluzioni la modellazione delle sequenze. Sviluppato da Albert Gu e Tri Dao, Mamba si distingue per la sua efficienza nell'elaborazione di sequenze complesse in campi come l'elaborazione del linguaggio, la genomica e l'analisi audio. La sua modellazione di sequenze in tempo lineare con spazi di stato selettivi garantisce prestazioni eccezionali in queste diverse modalità.
Approfondiremo la capacità di Mamba di superare le sfide computazionali dei Transformer tradizionali, soprattutto con sequenze lunghe. Il suo approccio selettivo nei modelli nello spazio di stato consente un'inferenza più rapida e un ridimensionamento lineare con la lunghezza della sequenza, migliorando significativamente la produttività.
L'unicità di Mamba risiede nella sua rapida capacità di elaborazione, nel livello SSM selettivo e nel design hardware-friendly ispirato a FlashAttention. Queste caratteristiche consentono a Mamba di superare molti modelli esistenti, inclusi quelli basati sull'approccio del trasformatore, rendendolo un progresso degno di nota nel machine learning.
Trasformatori contro Mamba
I trasformatori, come GPT-4, hanno stabilito nuovi standard nell'elaborazione del linguaggio naturale. Tuttavia, la loro efficienza diminuisce con sequenze più lunghe. È qui che Mamba fa un balzo in avanti, con la sua capacità di elaborare sequenze lunghe in modo più efficiente e la sua architettura unica che semplifica l'intero processo.
Trasformatori esperti nel gestire sequenze di dati, come testo per modelli linguistici. A differenza dei modelli precedenti che elaboravano i dati in sequenza, i Transformer elaborano intere sequenze simultaneamente, consentendo loro di acquisire relazioni complesse all'interno dei dati.
Usano il meccanismo dell'attenzione, che consente al modello di concentrarsi su diverse parti della sequenza quando fa 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 tutti gli altri elementi, fornendo un peso che indica l'importanza, o "attenzione", che ogni elemento dovrebbe ricevere quando si prevede l'elemento successivo nella sequenza.
I trasformatori mantengono due blocchi principali: l'encoder, che elabora i dati in ingresso, e il decoder, che genera l'output. Il codificatore è costituito da più strati, ciascuno contenente due sottostrati: un meccanismo di autoattenzione multi-testa e una semplice rete feed-forward completamente connessa in termini di posizione. La normalizzazione e le connessioni residue vengono utilizzate in ogni sottolivello per aiutare nell'addestramento delle reti profonde.
Il decoder presenta anche livelli con due sottolivelli simili a quelli dell'encoder, ma ne aggiunge un terzo che esegue l'attenzione multi-testa 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.
A differenza di Transformers, il modello Mamba adotta un approccio diverso. Mentre Transformers affronta il problema delle lunghe sequenze utilizzando meccanismi di attenzione più complessi, Mamba utilizza spazi di stato selettivi, fornendo una maggiore capacità di calcolo.
Ecco una panoramica generale del funzionamento di un trasformatore:
- Elaborazione dell'input: I trasformatori codificano innanzitutto i dati di input in un formato comprensibile dal modello, spesso utilizzando incorporamenti che incorporano anche la posizione di ciascun elemento nella sequenza.
- Meccanismo di attenzione: Fondamentalmente, il meccanismo dell'attenzione calcola un punteggio che rappresenta quanta attenzione mettere su altre parti della sequenza di input quando si comprende un elemento corrente.
- Architettura codificatore-decodificatore: Il modello del trasformatore è composto da un codificatore per elaborare l'input e da un decodificatore per generare l'output. Ciascuno è costituito da più livelli che perfezionano la comprensione dell'input da parte del modello.
- Attenzione multi-testa: Sia all'interno del codificatore che del decodificatore, l'attenzione multi-testa consente al modello di occuparsi simultaneamente di diverse parti della sequenza da diversi spazi rappresentazionali, migliorando la sua capacità di apprendere da contesti diversi.
- Reti feed-forward per posizione: Dopo l'attenzione, una semplice rete neurale elabora l'output di ciascuna posizione separatamente e in modo identico. Questo viene combinato con l'ingresso attraverso una connessione residua e seguito dalla normalizzazione dello strato.
- Generazione di output: Il decoder prevede quindi una sequenza di output, influenzata dal contesto del codificatore e da ciò che ha generato finora.
La capacità del trasformatore di gestire sequenze in parallelo e il suo robusto meccanismo di attenzione lo rendono potente per attività come la traduzione e la generazione di testo.
Al contrario, il modello Mamba opera in modo diverso, utilizzando spazi di stato selettivi per elaborare le sequenze. Questo approccio risolve l'inefficienza computazionale dei Transformers quando si gestiscono sequenze lunghe. Il design di Mamba consente un'inferenza più rapida e scala linearmente con la lunghezza della sequenza, definendo un nuovo paradigma per la modellazione delle sequenze che potrebbe essere più efficiente, soprattutto man mano che le sequenze diventano sempre più lunghe.
Mamba
Ciò che rende Mamba davvero unico è il suo allontanamento dall’attenzione tradizionale e dai blocchi MLP. Questa semplificazione porta a un modello più leggero e veloce che si adatta linearmente alla lunghezza della sequenza: un'impresa senza eguali rispetto ai suoi predecessori.
Le caratteristiche principali di Mamba includono:
- MVU selettivi: Questi consentono a Mamba di filtrare le informazioni irrilevanti e concentrarsi sui dati rilevanti, migliorando la gestione delle sequenze. Questa selettività è cruciale per un ragionamento efficiente basato sui contenuti.
- Algoritmo compatibile con l'hardware: Mamba utilizza un algoritmo parallelo ottimizzato per l'hardware moderno, in particolare per le GPU. Questo design consente un calcolo più veloce e riduce i requisiti di memoria rispetto ai modelli tradizionali.
- 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 prestazioni.
Mamba ha dimostrato prestazioni superiori in vari ambiti, tra cui linguaggio, audio e genomica, eccellendo sia in compiti di pre-addestramento che in compiti specifici di dominio. Ad esempio, nella modellazione linguistica, Mamba eguaglia o supera le prestazioni dei modelli Transformer più grandi.
Il codice di Mamba e i modelli pre-addestrati sono disponibili per l'uso da parte della comunità su GitHub.
Le attività di copia standard sono semplici per i modelli lineari. I responsabili della copia selettiva e dell'induzione richiedono una memoria dinamica e sensibile al contenuto per i LLM.
I modelli Structured State Space (S4) sono recentemente emersi come una classe promettente di modelli di sequenza, che comprendono tratti di RNN, CNN e modelli classici di spazio degli stati. 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 del deep learning, rappresentano un’innovazione significativa, fornendo una nuova metodologia per la progettazione di modelli di sequenza efficienti e altamente adattabili.
La dinamica dei modelli S4
MVS (S4) Questo è il modello base dello spazio degli stati strutturato. Ci vuole una sequenza x e produce un output y utilizzando i parametri appresi A, B, Ce un parametro di ritardo ΔLa trasformazione implica la discretizzazione dei parametri (trasformando le funzioni continue in discrete) e l'applicazione dell'operazione SSM, che è invariante nel tempo, ovvero non cambia nei diversi intervalli di tempo.
Il significato 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 un'adeguata normalizzazione, migliorando la stabilità e le prestazioni del modello. La discretizzazione traccia anche parallelismi con i meccanismi di controllo presenti nelle RNN, che sono fondamentali per la gestione del flusso di informazioni attraverso la rete.
Invarianza temporale lineare (LTI)
Una caratteristica fondamentale dei modelli S4 è la loro invarianza temporale lineare. Questa proprietà implica che la dinamica del modello rimanga coerente nel tempo, con i parametri fissati per tutti i passaggi temporali. LTI è una pietra angolare della ricorrenza e delle convoluzioni, offrendo un framework semplificato ma potente per la creazione di modelli di sequenza.
Superare i limiti fondamentali
Il framework S4 è stato tradizionalmente limitato dalla sua natura LTI, che pone sfide nella modellazione dei dati che richiedono dinamiche adattive. Il recente documento 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 una serie più diversificata di sequenze e attività, ampliando significativamente la loro applicabilità.
Il termine "modello nello spazio degli stati" copre in senso lato qualsiasi processo ricorrente che coinvolga uno stato latente ed è stato utilizzato per descrivere vari concetti in diverse discipline. Nel contesto del deep learning, i modelli S4, o SSM strutturati, si riferiscono a una classe specifica di modelli ottimizzati per un calcolo efficiente, pur mantenendo la capacità di modellare sequenze complesse.
I modelli S4 possono essere integrati in architetture di rete neurale end-to-end, funzionando come trasformazioni di sequenze autonome. Possono essere visti come analoghi agli strati di convoluzione nelle CNN, fornendo la spina dorsale per la modellazione di sequenze in una varietà di architetture di reti neurali.
MVU vs MVU + Selezione
Motivazione per la selettività nella modellazione di sequenze
MVU strutturati
L'articolo sostiene che un aspetto fondamentale della modellazione di sequenze è la compressione del contesto in uno stato gestibile. I modelli che possono concentrarsi selettivamente o filtrare gli input forniscono un mezzo più efficace per mantenere questo stato compresso, portando a modelli di sequenza più efficienti e potenti. Questa selettività è vitale affinché i modelli possano controllare in modo adattivo il modo in cui 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, il che introduce un grado di adattività precedentemente irraggiungibile con i modelli invarianti nel tempo. Ciò si traduce in SSM variabili nel tempo 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 la dipendenza dall'input ai parametri B e Ce un parametro di ritardo Δ. Ciò consente al modello di concentrarsi selettivamente su alcune 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 sequenza, regolando il focus dinamicamente nel tempo.
Punti salienti delle prestazioni di Mamba
Mamba è il migliore della categoria in ogni singolo risultato della valutazione
In termini di prestazioni, Mamba eccelle sia in termini di velocità di inferenza che di accuratezza. Il suo design consente un migliore utilizzo di contesti più lunghi, come dimostrato sia nella modellazione del DNA che in quella audio, superando i modelli precedenti in attività complesse che richiedono dipendenze a lungo raggio. La sua versatilità è evidenziata anche nelle valutazioni zero-shot su più attività, 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 prevede semplici comandi pip per installare i pacchetti necessari dal file Deposito Mamba. Se si verificano problemi di compatibilità con le versioni PyTorch, può essere utile utilizzare il flag –no-build-isolation con pip. Questi modelli, addestrati su set di dati estesi come il set di dati Pile e SlimPajama, sono progettati per soddisfare varie esigenze computazionali e benchmark di prestazioni.
Mamba offre diversi livelli di interfaccia, dal livello SSM selettivo al blocco Mamba e alle strutture complete del modello linguistico. Il blocco Mamba, che è il modulo principale dell'architettura, utilizza un livello Conv1d causale e può essere facilmente integrato nei progetti di reti neurali. L'esempio di utilizzo fornito in Python mostra come istanziare un modello Mamba e come elaborare i dati attraverso di esso, evidenziando la semplicità e la flessibilità del sistema.
I modelli Mamba preaddestrati sono disponibili su Abbracciare il viso, con dimensioni che vanno da 130 milioni a 2.8 miliardi di parametri, addestrati sull'ampio set di dati Pile e sul set di dati SlimPajama. Questi modelli sono progettati per soddisfare diversi requisiti computazionali e prestazionali, aderendo agli standard dimensionali di GPT-3. Gli utenti possono aspettarsi produttività e precisione elevate da questi modelli, rendendo Mamba una scelta competitiva per varie applicazioni, inclusa ma non limitata alla modellazione del linguaggio.
L'impatto di Mamba
Mamba rappresenta un balzo in avanti nella modellazione sequenziale, offrendo una potente alternativa alle architetture Transformer per l'elaborazione di dati ad alta densità 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 codice sorgente di Mamba e dei suoi modelli pre-addestrati lo rende uno strumento accessibile e robusto per ricercatori e sviluppatori nel campo dell'intelligenza artificiale e del deep learning.
