Lo sviluppo di Large Language Models (LLM) costruiti da modelli di trasformatori decoder-only ha giocato un ruolo cruciale nel trasformare il dominio del Natural Language Processing (NLP), nonché nell’avanzare diverse applicazioni di apprendimento automatico, tra cui reinforcement learning, analisi di serie temporali, elaborazione di immagini e molto altro. Tuttavia, nonostante la loro scalabilità e la forte performance, gli LLM costruiti da modelli di trasformatori decoder-only affrontano ancora limitazioni significative. Sebbene espressivi, il meccanismo di attenzione nei modelli LLM derivati dai trasformatori richiede risorse computazionali elevate durante sia l’inferenza che l’addestramento, necessitando di una sostanziale quantità di memoria per la lunghezza della sequenza e operazioni computazionali (FLOPs) che aumentano quadraticamente. Questa alta richiesta computazionale limita la lunghezza del contesto dei modelli, rende le attività di generazione autoregressiva proporzionalmente costose con la scala e ostacola la capacità dei modelli di apprendere da flussi di dati continui e di elaborare sequenze di lunghezza illimitata in modo efficiente.

Recentemente, Modelli di Spazio di Stato (SSM) hanno dimostrato capacità e prestazioni notevoli, competendo con modelli di architettura di trasformatori in benchmark di modellazione su larga scala, raggiungendo una complessità di memoria come funzione della lunghezza della sequenza e una complessità temporale lineare. Inoltre, Mamba, un recente modello SSM, ha mostrato prestazioni eccezionali in una gamma di attività di modellazione del linguaggio e di elaborazione di sequenze lunghe. Contestualmente, i modelli Mixture of Expert (MoE) hanno anche mostrato prestazioni impressionanti, riducendo significativamente la latenza e i costi computazionali dell’inferenza, sebbene a scapito di un’impronta di memoria più grande. Basandosi sui modelli Mamba e MoE, questo articolo discuterà di BlackMamba, una nuova architettura che combina il modello SSM Mamba con i modelli MoE per sfruttare i vantaggi offerti da entrambi i framework. Gli esperimenti su BlackMamba hanno dimostrato la sua capacità di superare il framework Mamba esistente e i modelli di trasformatori di base sia in termini di FLOPs di addestramento che di inferenza. La prestazione eccezionale del framework BlackMamba mostra che è in grado di combinare efficacemente le capacità dei framework Mamba e MoE, offrendo un’inferenza veloce e a basso costo da MoE con generazione a complessità lineare da Mamba.

Questo articolo si propone di coprire il framework BlackMamba in profondità. Esploreremo il meccanismo, la metodologia e l’architettura del framework, insieme al suo confronto con framework di generazione di immagini e video di stato dell’arte. Iniziamo.

BlackMamba: Introduzione a MoE per Modelli di Spazio di Stato

L’evoluzione dei Large Language Models (LLM), in particolare quelli basati su architetture di trasformatori decoder-only, ha notevolmente influenzato il campo del Natural Language Processing (NLP) e si è estesa a diverse applicazioni di apprendimento automatico, tra cui reinforcement learning, analisi di serie temporali, elaborazione di immagini e oltre. Tuttavia, nonostante la loro scalabilità e la forte performance, questi LLM basati su trasformatori decoder-only affrontano ancora sfide notevoli. Il meccanismo di attenzione, una caratteristica chiave dei modelli LLM basati su trasformatori, richiede risorse computazionali estese per l’inferenza e l’addestramento. Ciò comporta la necessità di memoria che cresce con la lunghezza della sequenza e operazioni computazionali (FLOPs) che aumentano quadraticamente. Tali esigenze computazionali intensive limitano la lunghezza del contesto dei modelli, aumentano i costi delle attività di generazione autoregressiva con la scala e ostacolano la capacità dei modelli di apprendere da flussi di dati continui o di elaborare sequenze di lunghezza illimitata in modo efficiente.

Sono stati fatti notevoli sforzi negli ultimi anni per superare queste limitazioni, e l’attenzione si è spostata verso l’elaborazione di alternative architettoniche ai modelli di trasformatori canonici con densità di attenzione, con i Modelli di Spazio di Stato (SSM) e i modelli Mixture of Expert (MoE) come i candidati architettonici più promettenti. Il principale vantaggio derivato dal favorire i Modelli di Spazio di Stato rispetto ai modelli di architettura di trasformatori è la complessità computazionale lineare rispetto alla lunghezza della sequenza di input offerta dagli SSM, in contrasto con la complessità quadratica offerta dai trasformatori. Teoricamente, la complessità computazionale lineare rispetto alla lunghezza della sequenza di input consente ai Modelli di Spazio di Stato di elaborare sequenze più lunghe rispetto ai modelli di architettura di trasformatori per un dato budget di FLOPS o operazioni in virgola mobile al secondo, e di rendere la generazione autoregressiva costante in termini di calcolo senza un cache KV. Recentemente sviluppati SSM, tra cui RWKV e Mamba, hanno dimostrato un’elaborazione efficiente di sequenze lunghe e un addestramento competitivo con i modelli di trasformatori, con proprietà di scalabilità simili. D’altra parte, i modelli Mixture of Expert stanno guadagnando popolarità come alternativa ai trasformatori densi, poiché facilitano una riduzione significativa dei FLOPs di inferenza e addestramento, essenziali per raggiungere una qualità paragonabile a un modello denso. I modelli MoE (Mixture of Experts) operano attivando solo una selezione sparsa dei parametri totali durante un singolo passo in avanti. Utilizzano una funzione di routing per determinare quali “esperti” sono attivati in base al contesto dato. Questo approccio crea una separazione tra il costo computazionale dell’inferenza e il numero totale di parametri, consentendo prestazioni migliorate all’interno di un budget di inferenza fisso, sebbene con un aumento del numero di parametri e una maggiore richiesta di memoria.

Questo avanzamento architettonico offre notevoli vantaggi rispetto ai trasformatori tradizionali e rappresenta una direzione emozionante per ulteriori sviluppi. Sosteniamo che integrare questi miglioramenti in un modello Mamba-MoE combinato potrebbe accelerare notevolmente le capacità di modellazione del linguaggio e l’efficienza oltre quella dei modelli di trasformatori standard. I vantaggi previsti di un’architettura Mamba-MoE rispetto a un modello di trasformatore denso tradizionale includono:

Mamba: raggiunge una complessità computazionale lineare rispetto alla lunghezza della sequenza di input per entrambe le fasi di addestramento e inferenza. Consente la generazione autoregressiva in un tempo costante e con utilizzo di memoria costante.

MoE: offre una velocità di inferenza e un’efficienza computazionale dell’addestramento paragonabili a un modello denso più piccolo, mantenendo un livello di qualità del modello che rivaluta quello di un modello con un numero equivalente di parametri rispetto alla versione più densa.

Tuttavia, è essenziale affermare che i modelli di architettura di trasformatori sono ancora lo stato dell’arte e hanno dimostrato prestazioni forti e coerenti in attività di modellazione del linguaggio e di elaborazione di sequenze. Al suo nucleo, l’architettura di trasformatori impiega un’attenzione che esegue un confronto all-to-all quadratico del prodotto dei punti tra le incapsulazioni di token diversi in una sequenza e esegue una mappatura lineare su un vettore di output. Il modello di trasformatori consiste in blocchi di auto-attenzione impilati tra blocchi MLP (Multi-Layer Perceptron) che ulteriormente consistono in un MLP a due strati con una funzione di attivazione data.

BlackMamba: Architettura e Metodologia

Modelli di Spazio di Stato

I Modelli di Spazio di Stato appartengono al gruppo di modelli di sequenza con complessità lineare rispetto alla lunghezza della sequenza di input. L’architettura dei Modelli di Spazio di Stato si allinea più con le Reti Neurali Ricorrenti e le Reti Neurali Convoluzionali piuttosto che con l’architettura basata sull’attenzione, e si ispira a un sistema dinamico continuo che mappa una funzione unidimensionale attraverso uno spazio latente implicito. Un sistema dinamico lineare rende efficienti i calcoli paralleli utilizzando una scansione associativa o una scansione convoluzionale. Nelle situazioni pratiche, la natura ricorrente dei Modelli di Spazio di Stato è stata la ragione per cui non è ancora stata adottata su hardware AI altamente parallelo come le GPU. Tuttavia, l’emergere di SSM come RWKV e Mamba hanno utilizzato kernel di scansione parallela per mappare efficientemente operazioni ricorrenti sulle GPU, facilitando così l’addestramento di nuove architetture con efficienza paragonabile a quella raggiunta dai modelli di trasformatori.

La complessità quadratica intrinseca rispetto alla lunghezza della sequenza all’interno dei trasformatori è una limitazione ben nota che ostacola il ragionamento e la comprensione su contesti molto lunghi. Recentemente, sono state introdotte innovazioni che hanno esteso la lunghezza del contesto, consentendo ai trasformatori di essere addestrati su una scala fattibile prima di essere applicati a contesti molto più lunghi durante l’inferenza. Nonostante questi progressi, il processo di inferenza richiede ancora una notevole quantità di risorse computazionali e di memoria, specialmente per mantenere la cache KV, rendendolo un’impresa intensiva in termini di risorse. Ricerche recenti si sono concentrate sull’aumentare le capacità espressive dei modelli di spazio di stato incorporando meccanismi di gating dipendenti dall’input, simili alle matrici QKV (Query, Key, Value) trovate nei meccanismi di attenzione.

Questi sforzi mirano a preservare la progressione lineare intrinseca della ricorsione dello spazio di stato, consentendo un’esecuzione efficiente attraverso una scansione convoluzionale o selettiva. Questo approccio riduce notevolmente il divario di prestazioni con i trasformatori nelle applicazioni pratiche. Tra questi progressi, Mamba si distingue come un modello di spazio di stato che riflette gli obiettivi della ricerca precedente, mostrando livelli di prestazione impressionanti paragonabili a quelli dei trasformatori su scale fino a 2,8 miliardi di parametri. Ciò viene raggiunto applicando un gating dipendente dall’input agli input della ricorsione del modello SSM, assicurando al contempo calcoli efficienti attraverso l’uso di kernel di scansione selettiva personalizzati.

Modelli Mixture of Expert

I modelli Mixture of Expert (MoE) raggiungono una separazione tra il costo dell’inferenza e il numero totale di parametri attivando selettivamente i parametri durante il passo in avanti. Invece di utilizzare tutti i parametri, questi modelli dirigono i token verso specifici esperti MLP (Multilayer Perceptron). Idealmente, ogni esperto è progettato per elaborare un tipo specifico di input, con un meccanismo di routing, essenzialmente una rete neurale compatta, che determina l’esperto più adatto per ogni token. Questo approccio mira a preservare la potenza espressiva complessiva di un modello con un numero equivalente di parametri in una configurazione più densa, ma con esigenze computazionali notevolmente ridotte. Tipicamente, il router è una mappatura dei layer lineari dai token agli indici degli esperti, con ogni esperto che è semplicemente un MLP standard. Tuttavia, gli sviluppatori devono ancora determinare il metodo di addestramento ottimale per il router, poiché il problema di assegnazione degli esperti non è differenziabile, e i modelli MoE spesso lottano con il bilanciamento del carico e la stabilità dell’addestramento tra diversi esperti per l’efficienza hardware.

Architettura

Al suo nucleo, BlackMamba impiega un modello di trasformatore standard composto da blocchi MLP e blocchi di attenzione interpolati lungo un flusso residuale. Ora, la maggior parte dei modelli MoE sostituisce semplicemente i blocchi MLP con un layer di esperti instradati. D’altra parte, il framework BlackMamba non solo sostituisce il blocco MLP nel trasformatore con un layer di esperti instradati, ma sostituisce anche il blocco di attenzione con un layer del modello SSM Mamba. L’architettura del framework BlackMamba è illustrata nella figura seguente.

Addestramento e Dataset

Il modello BlackMamba è addestrato su oltre 300 miliardi di token su un dataset personalizzato e utilizza la funzione di attivazione SwiGLU per i percettroni multistrato degli esperti. Il framework si addestra con 8 esperti, un numero che gli sviluppatori hanno trovato essere il giusto equilibrio e compromesso tra l’impronta di memoria e il costo dell’inferenza del modello. Il dataset personalizzato utilizzato per addestrare il framework BlackMamba consiste in una miscela di dataset open source esistenti, tra cui Starcoder, SlimPajama, Pile e altri. La tabella seguente illustra i pesi di ciascun dataset utilizzato per l’addestramento del framework BlackMamba. Complessivamente, ci sono 1,8 trilioni di token nel dataset.

BlackMamba: Risultati

Per garantire un confronto equo tra Mamba e BlackMamba, gli sviluppatori hanno addestrato entrambi i modelli con gli stessi parametri di addestramento sui dati di addestramento. Il framework BlackMamba è in grado di superare sia il framework Mamba esistente che i modelli di trasformatori di base sia in termini di FLOPs di addestramento che di inferenza. La figura seguente illustra il tempo necessario per generare una sequenza di una data lunghezza in modo autoregressivo da un prompt iniziale di un token come funzione della lunghezza della sequenza.

Inoltre, i vantaggi di latenza sia dei modelli MoE che di Mamba sono combinati nel framework BlackMamba, risultando in tempi di inferenza significativamente più veloci rispetto ai modelli di trasformatori, ai modelli Mamba puri e ai modelli MoE. Inoltre, il vantaggio dell’inferenza del framework BlackMamba è direttamente proporzionale alle lunghezze delle sequenze, rendendo BlackMamba estremamente efficace nella generazione di sequenze lunghe. Proseguendo, la figura seguente illustra il numero di token assegnati ai modelli BlackMamba con 340 milioni e 640 milioni di parametri rispettivamente. Come si può vedere, la maggior parte dei layer mostra un alto livello di equilibrio degli esperti a causa dell’algoritmo Sinkhorn migliorato implementato dai modelli BlackMamba.

La tabella seguente copre i punteggi di valutazione del framework BlackMamba in confronto a una gamma di modelli linguistici pre-addestrati open source. Come si può osservare, il framework BlackMamba è in grado di competere e superare la maggior parte dei framework su tutte le basi di confronto. Inoltre, è degno di nota che i modelli che superano BlackMamba hanno un numero di parametri notevolmente più alto, e il divario di prestazioni è minimo, indicando la capacità del framework BlackMamba di raggiungere prestazioni elevate con meno parametri.

Pensieri Finali

In questo articolo, abbiamo discusso di BlackMamba, una nuova architettura che combina il modello SSM Mamba con i modelli MoE per sfruttare i vantaggi offerti da entrambi i framework. Gli esperimenti su BlackMamba hanno dimostrato la sua capacità di superare il framework Mamba esistente e i modelli di trasformatori di base sia in termini di FLOPs di addestramento che di inferenza. La prestazione eccezionale del framework BlackMamba mostra che è in grado di combinare efficacemente le capacità dei framework Mamba e MoE, offrendo un’inferenza veloce e a basso costo da MoE con generazione a complessità lineare da Mamba. Abbiamo discusso di come l’architettura del framework BlackMamba sia in grado di superare i modelli di trasformatori addestrati con forza, il framework Mamba esistente e i modelli MoE in termini di FLOPs di addestramento e costo di inferenza. Inoltre, il framework BlackMamba eredita anche i FLOPs di generazione e l’addestramento ridotto da entrambi i modelli MoE e Mamba contemporaneamente.