Introduzione all’architettura Mamba LLM: un nuovo paradigma nel machine learning

I modelli linguistici sono un tipo di modello di machine learning addestrato a definire una distribuzione di probabilità sul linguaggio naturale. La loro architettura è composta principalmente da più livelli di reti neurali, come livelli ricorrenti, feedforward, di embedding e di attenzione. Questi livelli si combinano per elaborare un testo in input e generare previsioni in output.

Alla fine del 2023, ricercatori della Carnegie Mellon e della Princeton University hanno pubblicato un paper di ricerca che presentava una nuova architettura per i Large Language Models (LLM) chiamata Mamba. Mamba è una nuova architettura di modelli a spazio di stato incentrata sul sequence modeling. È stata sviluppata per superare alcune limitazioni dei modelli transformer, soprattutto nell’elaborazione di sequenze lunghe, e sta mostrando prestazioni promettenti.

Esploriamo l’architettura di Mamba LLM e la sua importanza nel machine learning.

Che cos’è Mamba?

Mamba è una nuova architettura LLM che integra il modello S4 (Structured State Space) per gestire sequenze di dati lunghe. Combinando le migliori caratteristiche dei modelli ricorrenti, convoluzionali e a tempo continuo, S4 riesce a simulare in modo efficace ed efficiente le dipendenze a lungo termine. Questo gli consente di gestire dati campionati in modo irregolare, avere un contesto non limitato e mantenere l’efficienza computazionale durante addestramento e test.

Espandendo il paradigma S4, Mamba introduce diversi miglioramenti notevoli, soprattutto nella gestione di operazioni variabili nel tempo. La sua architettura ruota attorno a un particolare meccanismo di selezione che modifica i parametri del modello a spazio di stato (SSM) in base all’input.

Di conseguenza, Mamba può filtrare con successo i dati meno importanti concentrandosi solo sulle informazioni cruciali all’interno delle sequenze. Secondo Wikipedia, "Il modello passa da un quadro invariabile nel tempo a uno variabile nel tempo, il che influisce sia sul calcolo sia sull’efficienza del sistema."

Caratteristiche e innovazioni principali

La deviazione di Mamba dall’attenzione convenzionale e dai blocchi MLP lo distingue. Questa semplificazione produce un modello più leggero, più veloce e che scala linearmente con la lunghezza della sequenza, un risultato mai raggiunto dai suoi predecessori.

I componenti chiave di Mamba includono:

• Selective-State-Spaces (SSM): I modelli ricorrenti, che elaborano le informazioni in modo selettivo a seconda dell’input corrente, sono il fondamento degli SSM di Mamba. Questo consente di filtrare i dati superflui e concentrarsi su quelli pertinenti, con un potenziale aumento dell’efficienza di elaborazione.
• Architettura semplificata: Mamba sostituisce i complessi blocchi di attenzione e MLP dei Transformer con un unico blocco SSM coeso. L’obiettivo è accelerare l’inferenza e ridurre la complessità computazionale.
• Parallelismo ottimizzato per l’hardware: Mamba può ottenere prestazioni ancora migliori grazie a una modalità ricorrente con un algoritmo parallelo ideato specificamente per l’efficienza hardware.

Un altro componente chiave è l’invarianza nel tempo lineare (LTI); la LTI è una delle caratteristiche fondamentali dei modelli S4. Questa proprietà suggerisce che i parametri del modello restano costanti in tutti i timestep, mantenendo la coerenza nella dinamica del modello. Costruire modelli di sequenza è più semplice ed efficace con la LTI, che è alla base di ricorrenza e convoluzioni.

Dettagli dell’architettura Mamba LLM

L’architettura di Mamba sottolinea ulteriormente l’importanza dei progressi compiuti nel machine learning. Modifica il modo in cui i modelli elaborano le sequenze introducendo un livello SSM selettivo. Questo permette a Mamba di fare due cose estremamente importanti:

1. Concentrarsi sulle informazioni rilevanti – Mamba può dare priorità ai dati più predittivi per il compito assegnando un peso diverso a ciascun input.

2. Adattarsi dinamicamente agli input – Grazie alla capacità del modello di adattarsi all’input, Mamba può gestire con facilità vari compiti di sequence modeling.

Di conseguenza, Mamba può gestire sequenze con un’efficienza senza precedenti, il che lo rende la scelta ideale per compiti che coinvolgono lunghe sequenze di dati.

La filosofia progettuale di Mamba si basa sulla consapevolezza delle capacità dell’hardware moderno. È progettato per sfruttare appieno la potenza di calcolo delle GPU, garantendo:

• Uso della memoria ottimizzato: La progettazione dell’espansione dello stato di Mamba per rientrare nella memoria HBM (high-bandwidth memory) delle GPU riduce i tempi di trasferimento dei dati e accelera l’elaborazione.

• Massimizzazione dell’elaborazione parallela: Allineando i propri calcoli con la natura parallela del computing su GPU, Mamba raggiunge un livello di prestazioni che stabilisce un nuovo standard per i modelli di sequenza.

Mamba vs Transformer

L’introduzione dei Transformer, come GPT-4, ha rivoluzionato il campo dell’elaborazione del linguaggio naturale (NLP) e fissato benchmark per diversi compiti di linguaggio naturale. Le sequenze più lunghe sono da tempo un punto dolente per i Transformer, perché ne compromettono sensibilmente l’efficienza.

È proprio qui che Mamba eccelle. In particolare, Mamba può elaborare sequenze lunghe più rapidamente dei Transformer e lo fa in modo più semplice grazie alla sua architettura unica.

Architettura dei Transformer

I Transformer sono estremamente abili nel lavorare con sequenze di dati, come il testo per i modelli linguistici. Elaborano intere sequenze contemporaneamente, a differenza dei modelli precedenti che elaboravano i dati in modo sequenziale. Questa caratteristica intrinseca consente loro di cogliere relazioni complesse all’interno dei dati.

Usano un meccanismo di attenzione che permette al modello di concentrarsi su diverse porzioni della sequenza durante la generazione delle previsioni. Per calcolare questa attenzione si usano tre insiemi di pesi: values, keys e queries ottenuti dai dati in input.

Ogni elemento in una sequenza viene pesato rispetto a tutti gli altri per indicare quanta importanza — o "attenzione" — dovrebbe avere nel prevedere l’elemento successivo della serie.

I Transformer comprendono due blocchi principali: il decoder, che crea l’output, e l’encoder, che elabora i dati in input.

L’encoder è composto da diversi strati — ciascuno con due sotto-strati: una rete feed-forward fully connected, position-wise, di base e un meccanismo di self-attention multi-head. Per facilitare l’addestramento di reti profonde, a ogni sotto-strato si applicano connessioni residue e normalizzazione.

Come l’encoder, il decoder è composto da due strati con due sotto-strati, ma aggiunge anche un terzo sotto-strato che gestisce la multi-head attention sull’output dell’encoder. La proprietà autoregressiva del decoder è preservata grazie alla sua natura sequenziale, che limita le previsioni per una posizione a considerare solo le posizioni precedenti.

Dunque, i Transformer cercano di affrontare il problema delle sequenze lunghe utilizzando meccanismi di attenzione più complessi, ma Mamba adotta un approccio diverso.

Architettura di Mamba

Mamba sfrutta spazi di stato selettivi. Questo metodo risolve le inefficienze computazionali dei Transformer con sequenze lunghe.

L’architettura di Mamba consente un’inferenza più rapida e una scalabilità lineare con la lunghezza della sequenza, creando un nuovo paradigma per il sequence modeling che può rivelarsi più efficace man mano che le sequenze si allungano.

Dato che abbiamo già approfondito l’architettura di Mamba sopra, qui non ci addentreremo oltre.

Ecco un grafico da Wikipedia per visualizzare meglio il confronto tra Mamba e i Transformer:

Vale la pena notare che, nonostante i numerosi vantaggi che gli SSM hanno rispetto ai Transformer, questi ultimi possono gestire sequenze molto più lunghe di quelle che gli SSM possono memorizzare, richiedono molti meno dati per apprendere compiti simili e superano gli SSM in compiti che richiedono recupero o copia dal contesto di input, anche con meno parametri. 

Primi passi con Mamba

Se vuoi sperimentare con Mamba o usarlo in un progetto, ti serviranno:

• Linux
• GPU NVIDIA
• PyTorch 1.12+
• CUDA 11.6+

Per installare i pacchetti necessari dal repository Mamba, usa alcune semplici istruzioni pip:

• [Opzionale] pip install causal-conv1d>=1.2.0: un’implementazione efficiente di un semplice livello Conv1d causale usato all’interno del blocco Mamba.
• pip install mamba-ssm: il pacchetto Mamba principale.

Può anche essere compilato dai sorgenti con pip install . da questo repository.

Se le versioni di PyTorch creano problemi di compatibilità, si può usare pip con l’opzione --no-build-isolation per aiutare. Questi modelli sono stati addestrati su grandi dataset come Pile e SlimPajama e sono stati costruiti per soddisfare vari requisiti di calcolo e benchmark prestazionali.

Il modello Mamba ha diversi livelli di interfaccia, ma il modulo principale è il blocco architetturale Mamba che incapsula l’SSM selettivo.

Una volta installato tutto, lo si può usare così:

# Source: Mamba Repository
import torch
from mamba_ssm import Mamba

batch, length, dim = 2, 64, 16
x = torch.randn(batch, length, dim).to("cuda")
model = Mamba(
    # This module uses roughly 3 * expand * d_model^2 parameters
    d_model=dim, # Model dimension d_model
    d_state=16,  # SSM state expansion factor
    d_conv=4,    # Local convolution width
    expand=2,    # Block expansion factor
).to("cuda")
y = model(x)
assert y.shape == x.shape

Applicazioni di Mamba

L’introduzione di Mamba LLM rappresenta un potenziale cambiamento significativo nello spazio delle architetture LLM. Più veloce, efficiente e scalabile, Mamba gestisce con facilità sequenze lunghe mantenendo standard prestazionali elevati, motivo per cui è destinato a svolgere un ruolo cruciale nel plasmare il futuro dei sistemi di IA avanzati.

In particolare, l’efficacia e le prestazioni di Mamba potrebbero dare origine alla prossima ondata di innovazioni nell’IA, aprendo la strada alla creazione di modelli e applicazioni sempre più complessi. Il suo potenziale impatto è enorme e comprende applicazioni di elaborazione audio e vocale, analisi di testi lunghi, creazione di contenuti, traduzione linguistica in tempo reale e altro ancora.

I settori che potrebbero essere rivoluzionati includono:

• Sanità: Mamba potrebbe accelerare lo sviluppo di cure personalizzate analizzando rapidamente i dati genetici.

• Finanza: Mamba potrebbe essere impiegato per analizzare tendenze di mercato di lungo periodo, con previsioni sui titoli più accurate.

• Assistenza clienti: Mamba potrebbe alimentare chatbot che tengono traccia di conversazioni estese, migliorando le comunicazioni con i clienti.

La strada avanti per Mamba

Mamba si distingue come un faro di innovazione che indica nuove opportunità per affrontare problemi complessi di sequence modeling. Il suo debutto è un passo verso sistemi più intelligenti, scalabili ed efficienti, in grado di comprendere e gestire sequenze con una profondità mai vista prima.

Sebbene Mamba si sia dimostrato un traguardo tecnico significativo, il suo successo non dipende solo dalle capacità tecnologiche — ricerca collaborativa, risorse condivise e contributi open source svolgono tutti un ruolo essenziale:

• Contributi open source: Incoraggiare ricercatori e sviluppatori a contribuire al codebase di Mamba può portare a modelli più resilienti e adattabili.

• Risorse condivise: La comunità può accelerare i progressi unendo conoscenze e risorse e condividendo modelli pre-addestrati.

• Ricerca collaborativa: Collaborazioni tra istituzioni accademiche e aziende possono ampliare le capacità di Mamba.

Conclusione

Mamba non offre solo un miglioramento incrementale rispetto agli attuali modelli di sequenza; ridefinisce ciò che è possibile. Con la sua introduzione, la storia dell’intelligenza artificiale vedrà un nuovo capitolo in cui le inefficienze computazionali e i limiti di lunghezza delle sequenze stanno finalmente diventando obsoleti.

Negli ultimi anni abbiamo visto l’evoluzione dell’IA dagli RNN ai Transformer e ora a Mamba, ogni passo più vicino a realizzare un’IA capace di un’elaborazione profonda e di un trattamento delle informazioni paragonabile a quello umano. Mamba incarna lo spirito rivoluzionario che spinge avanti il campo dell’IA con il suo approccio a spazio di stato selettivo e la scalabilità in tempo lineare.

Mamba segna l’inizio di uno sviluppo promettente nell’intelligenza artificiale. È un paradigma pensato per il futuro e destinato a incidere significativamente sull’IA grazie al suo potenziale illimitato.

Per continuare a imparare, dai un’occhiata a:

• Come creare applicazioni LLM con LangChain
• Come addestrare un LLM con PyTorch