Lo sviluppo di Large Language Models (LLM) costruiti a partire da modelli di trasformatori solo decodificatori ha svolto un ruolo cruciale nella trasformazione del dominio dell’elaborazione del linguaggio naturale (NLP), oltre a far avanzare diverse applicazioni di deep learning tra cui apprendimento per rinforzo , analisi di serie temporali, immagini elaborazione e molto altro ancora. Tuttavia, nonostante la loro scalabilità e le ottime prestazioni, gli LLM costruiti a partire da modelli di trasformatori basati solo su decoder devono ancora affrontare carenze significative. Sebbene espressivo, il meccanismo di attenzione nei LLM derivati dal trasformatore richiede elevate risorse computazionali sia durante l’inferenza che durante l’addestramento, richiedendo una memoria sostanziale per la lunghezza della sequenza e i FLOP quadratici. Questo elevato requisito computazionale limita la lunghezza del contesto dei modelli del trasformatore, rendendo le attività di generazione autoregressiva proporzionalmente costose con la scala e ostacola l’apprendimento da flussi di dati continui e la capacità di elaborazione di sequenze veramente illimitate.
In tempi recenti, i modelli State Space (SSM) hanno dimostrato capacità e prestazioni notevoli, competendo con i modelli di architettura del trasformatore in benchmark di modellazione su larga scala e raggiungendo al contempo la complessità della memoria in funzione della lunghezza della sequenza e del tempo lineare. Inoltre, Mamba, un modello spaziale dello stato recentemente pubblicato, ha mostrato prestazioni eccezionali in una serie di attività di modellazione del linguaggio e di elaborazione di lunghe sequenze. Allo stesso tempo, anche i modelli Mixture of Expert (MoE) hanno mostrato prestazioni impressionanti riducendo significativamente la latenza e i costi computazionali dell’inferenza, anche se a scapito di un ingombro di memoria maggiore. Basandosi sui modelli Mamba e MoE, questo articolo discuterà di BlackMamba, una nuova architettura che combina il modello Mamba State Space con i modelli MoE per sfruttare i vantaggi offerti da entrambi i framework. Gli esperimenti su BlackMamba hanno dimostrato la sua capacità di sovraperformare il framework Mamba esistente e le linee di base del trasformatore sia nei FLOP di addestramento che nell’inferenza. Le prestazioni eccezionali del framework BlackMamba dimostrano che può combinare efficacemente le capacità dei framework Mamba e MoE, offrendo un’inferenza rapida ed economica da MoE con la generazione di complessità lineare da Mamba.
La progressione dei Large Language Models (LLM), in particolare quelli basati su architetture di trasformatori basati solo su decoder, ha influenzato notevolmente il campo dell’elaborazione del linguaggio naturale (NLP) e si è espanso in varie applicazioni di deep learning, tra cui apprendimento per rinforzo, analisi di serie temporali, elaborazione di immagini , e oltre. Tuttavia, nonostante la loro scalabilità e le robuste prestazioni, questi LLM basati su trasformatore e solo decoder incontrano notevoli sfide. Il meccanismo di attenzione, una caratteristica chiave dei LLM basati su trasformatori , richiede ampie risorse computazionali sia per l’inferenza che per l’addestramento. Ciò comporta una necessità di memoria che cresce con la lunghezza della sequenza e le operazioni di calcolo (FLOP) che aumentano quadraticamente. Tali intense esigenze computazionali limitano la lunghezza del contesto dei modelli, aumentano i costi delle attività di generazione autoregressiva man mano che il modello si ridimensiona e ostacolano la capacità dei modelli di apprendere da flussi di dati continui o di elaborare sequenze di lunghezza illimitata in modo efficiente.
Negli ultimi anni sono stati compiuti sforzi significativi nel tentativo di superare queste limitazioni e l’attenzione è stata spostata verso l’ideazione di alternative architettoniche ai modelli canonici di trasformatore di attenzione denso, con i modelli SSM e MoE che sono le architetture candidate più promettenti. Il vantaggio principale ottenuto privilegiando i modelli di spazio degli stati rispetto ai modelli di architettura dei trasformatori è la complessità computazionale lineare rispetto alla lunghezza della sequenza di input offerta dagli SSM rispetto alla complessità quadratica offerta dai trasformatori. Teoricamente, la complessità computazionale lineare rispetto alla lunghezza della sequenza di input consente ai modelli dello spazio degli stati di elaborare sequenze più grandi rispetto ai modelli di architettura del trasformatore per un dato budget FLOPS o operazioni in virgola mobile al secondo e di rendere costante la generazione autoregressiva nel calcolo senza una cache KV. I modelli State Space sviluppati di recente, tra cui Mamba, RetNet e alcuni altri, hanno dimostrato un’efficiente inferenza e addestramento a lunga sequenza, insieme a prestazioni di attività di modellazione linguistica competitiva per trasformatori con proprietà di ridimensionamento simili. D’altra parte, le architetture di modelli Mixture of Expert stanno guadagnando popolarità come alternativa ai trasformatori densi poiché facilitano una significativa riduzione dell’inferenza e dei FLOP di addestramento essenziali per ottenere una qualità paragonabile a un modello denso. I modelli MoE (Mixture of Experts) funzionano attivando solo una selezione sparsa dei parametri totali durante un singolo passaggio in avanti. Utilizzano una funzione di routing per determinare quali “esperti” sono chiamati in azione 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, anche se con un numero maggiore di parametri e un requisito di memoria maggiore.
Questo progresso nell’architettura offre notevoli vantaggi rispetto ai trasformatori tradizionali e rappresenta una direzione entusiasmante per ulteriori sviluppi. Riteniamo che l’integrazione di questi miglioramenti in un modello combinato Mamba-MoE potrebbe accelerare significativamente le capacità e l’efficienza della modellazione del linguaggio oltre quella dei modelli di trasformazione standard. I vantaggi previsti di un’architettura Mamba-MoE rispetto a un modello tradizionale di trasformatore denso includono:
Mamba: raggiunge una complessità computazionale lineare relativa alla lunghezza della sequenza di input sia per le fasi di training che per quelle di inferenza. Consente la generazione autoregressiva in un intervallo di tempo costante e con un utilizzo costante della memoria.
MoE: offre velocità di inferenza ed efficienza computazionale di training paragonabili a un modello di base più piccolo e denso, pur mantenendo un livello di qualità del modello che rivaleggia con quello di un modello con un numero equivalente di parametri della versione più densa.
Detto questo, è essenziale affermare che i modelli di architettura dei trasformatori sono ancora all’avanguardia e hanno dimostrato prestazioni consistenti e notevoli nelle attività di modellazione del linguaggio e di elaborazione delle sequenze. Fondamentalmente, l’architettura del trasformatore impiega l’autoattenzione che esegue un confronto quadratico tutto a tutti delle somiglianze del prodotto scalare tra gli incorporamenti di diversi token in una sequenza ed esegue una mappa lineare su un vettore di output. Il modello del trasformatore è costituito da blocchi di auto-attenzione impilati tra blocchi MLP o Perceptron multistrato che consistono inoltre in un MLP a due strati con una determinata funzione di attivazione.
I modelli dello spazio degli stati appartengono al gruppo dei modelli di sequenza con complessità lineare rispetto alla lunghezza della sequenza di input. L’architettura dei modelli statali spaziali si allinea maggiormente 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 di convoluzione. Negli scenari pratici, la natura ricorrente dei modelli statali spaziali è stata la ragione per cui devono ancora essere adottati su hardware AI altamente paralleli come le GPU. Tuttavia, l’emergere di SSM come RWKV e Mamba hanno utilizzato kernel di scansione parallela per mappare in modo efficiente le operazioni ricorrenti sulle GPU, facilitando così l’addestramento di nuove architetture con efficienza paragonabile a quella ottenuta dai modelli di trasformazione.
La complessità quadratica intrinseca in relazione alla lunghezza della sequenza all’interno dei trasformatori è una limitazione ben nota che impedisce il ragionamento e la comprensione su contesti molto lunghi. Recenti innovazioni hanno introdotto l’idea di estendere 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 di calcolo e memoria, in particolare per il mantenimento della cache dei valori-chiave (KV), rendendolo un’attività ad alta intensità di risorse. Recenti sforzi di ricerca si sono concentrati sul miglioramento delle capacità espressive dei modelli dello spazio degli stati incorporando meccanismi di gating dipendenti dall’input, simili alle matrici Query, Key, Value (QKV) trovate nei meccanismi di attenzione.
Questi sforzi mirano a preservare la progressione intrinsecamente lineare della ricorsione nello spazio degli stati, consentendo un’esecuzione efficiente attraverso la convoluzione o un processo di scansione selettiva. Questo approccio riduce significativamente la disparità di prestazioni rispetto ai trasformatori nelle applicazioni pratiche. Tra questi progressi, Mamba si distingue come modello nello spazio degli stati che rispecchia gli obiettivi della ricerca precedente, mostrando livelli di prestazioni impressionanti paragonabili a trasformatori su scale fino a 2,8 miliardi di parametri. Ciò si ottiene applicando un gating dipendente dall’input agli input della ricorsione del modello nello spazio degli stati (SSM), garantendo nel contempo un calcolo efficiente attraverso l’uso di kernel di scansione selettiva su misura.
I modelli Mixture of Expert (MoE) raggiungono una separazione tra il costo di inferenza e il conteggio totale dei parametri attivando selettivamente i parametri durante il passaggio in avanti. Invece di utilizzare tutti i parametri, questi modelli indirizzano i token a specifici esperti di Multilayer Perceptron (MLP). Idealmente, ogni esperto è adattato per elaborare un particolare tipo di input, con un meccanismo di routing, essenzialmente una rete neurale compatta, che determina l’esperto più adatto per ciascun 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 richieste computazionali notevolmente ridotte. In genere, il router è una mappatura degli strati lineari dai token agli indici esperti in cui ciascun esperto è semplicemente un trasformatore standard Perceptron multistrato. Tuttavia, gli sviluppatori devono ancora individuare il metodo di addestramento ottimale per il router poiché il problema dell’assegnazione degli esperti non è differenziabile e i modelli di combinazione di esperti spesso hanno difficoltà con il bilanciamento del carico e la stabilità dell’addestramento tra diversi esperti per l’efficienza dell’hardware.
Fondamentalmente, BlackMamba utilizza un modello di trasformatore standard costituito da blocchi MLP interlacciati e blocchi di attenzione aggiunti in sequenza lungo un flusso residuo. Ora, la maggior parte dei modelli Mixture of Expert sostituisce semplicemente i blocchi di percettrone multistrato con uno strato esperto instradato. D’altra parte, il framework BlackMamba non solo sostituisce il blocco multistrato di percettrone nel trasformatore con uno strato esperto instradato, ma sostituisce anche lo strato di attenzione con uno strato di modello spaziale dello stato di Mamba. L’architettura del framework BlackMamba è mostrata nella figura seguente.
Il modello BlackMamba viene addestrato su oltre 300 miliardi di token su un set di dati personalizzato e utilizza la funzione di attivazione SwiGLU per i perceptron multistrato esperti. Il framework viene addestrato con 8 esperti, un numero che gli sviluppatori hanno ritenuto essere il giusto equilibrio e compromesso tra l’impronta di memoria e il costo di inferenza del modello. Il set di dati personalizzato utilizzato per addestrare il framework BlackMamba è costituito da una combinazione di set di dati open source già esistenti tra cui Starcoder, SlimPajama, Pile e altri.