L’ascesa di Mixture-of-Experts per modelli linguistici efficienti e di grandi dimensioni

Nel mondo dell’elaborazione del linguaggio naturale (NLP), la ricerca di costruire modelli linguistici più grandi e più capaci è stata una forza trainante dietro molti recenti progressi. Tuttavia, poiché questi modelli crescono in dimensioni, le esigenze computazionali per l’addestramento e l’inferenza diventano sempre più impegnative, spingendo contro i limiti delle risorse hardware disponibili.

Entra Mixture-of-Experts (MoE), una tecnica che promette di alleviare questo onere computazionale consentendo l’addestramento di modelli linguistici più grandi e più potenti. Di seguito, discuteremo MoE, esploreremo le sue origini, il suo funzionamento interno e le sue applicazioni nei modelli linguistici basati su transformer.

Le origini di Mixture-of-Experts

Il concetto di Mixture-of-Experts (MoE) può essere fatto risalire ai primi anni ’90, quando i ricercatori hanno esplorato l’idea di calcolo condizionale, in cui parti di una rete neurale sono attivate selettivamente in base ai dati di input. Uno dei lavori pionieristici in questo campo è stato il paper “Adaptive Mixture of Local Experts” di Jacobs et al. nel 1991, che ha proposto un framework di apprendimento supervisionato per un ensemble di reti neurali, ciascuna specializzata in una diversa regione dello spazio di input.

L’idea fondamentale dietro MoE è quella di avere più “esperti” reti, ciascuna responsabile dell’elaborazione di un subset dei dati di input. Un meccanismo di gating, tipicamente una rete neurale essa stessa, determina quali esperti dovrebbero elaborare un dato input. Questo approccio consente al modello di allocare le sue risorse computazionali in modo più efficiente attivando solo gli esperti rilevanti per ogni input, anziché impiegare la piena capacità del modello per ogni input.

Nel corso degli anni, vari ricercatori hanno esplorato e esteso l’idea di calcolo condizionale, portando a sviluppi come MoE gerarchici, approssimazioni a basso rango per il calcolo condizionale e tecniche per la stima dei gradienti attraverso neuroni stocastici e funzioni di attivazione a soglia rigida.

Mixture-of-Experts nei transformer

Mixture of Experts

Sebbene l’idea di MoE sia stata intorno per decenni, la sua applicazione ai modelli linguistici basati su transformer è relativamente recente. I transformer, che sono diventati lo standard de facto per i modelli linguistici di stato dell’arte, sono composti da più livelli, ciascuno contenente un meccanismo di auto-attenzione e una rete neurale feed-forward (FFN).

L’innovazione chiave nell’applicazione di MoE ai transformer è sostituire i livelli FFN densi con livelli MoE sparsi, ciascuno consistente in più esperti FFN e un meccanismo di gating. Il meccanismo di gating determina quali esperti dovrebbero elaborare ogni token di input, abilitando il modello a selezionare solo un subset di esperti per una data sequenza di input.

Uno dei primi lavori che ha dimostrato il potenziale di MoE nei transformer è stato il paper “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer” di Shazeer et al. nel 2017. Questo lavoro ha introdotto il concetto di un livello MoE con gating sparsamente, che ha impiegato un meccanismo di gating che ha aggiunto sparsità e rumore al processo di selezione degli esperti, garantendo che solo un subset di esperti fosse attivato per ogni input.

Da allora, molti altri lavori hanno ulteriormente avanzato l’applicazione di MoE ai transformer, affrontando sfide come l’instabilità di addestramento, il bilanciamento del carico e l’inferenza efficiente. Esempi notevoli includono il Switch Transformer (Fedus et al., 2021), ST-MoE (Zoph et al., 2022) e GLaM (Du et al., 2022).

Vantaggi di Mixture-of-Experts per i modelli linguistici

Il beneficio principale dell’impiego di MoE nei modelli linguistici è la capacità di aumentare le dimensioni del modello mantenendo un costo computazionale relativamente costante durante l’inferenza. Selezionando selettivamente solo un subset di esperti per ogni token di input, i modelli MoE possono raggiungere la potenza espressiva di modelli densi molto più grandi mentre richiedono significativamente meno calcoli.

Ad esempio, consideriamo un modello linguistico con un livello FFN denso di 7 miliardi di parametri. Se sostituiamo questo livello con un livello MoE composto da otto esperti, ciascuno con 7 miliardi di parametri, il numero totale di parametri aumenta a 56 miliardi. Tuttavia, durante l’inferenza, se attiviamo solo due esperti per token, il costo computazionale è equivalente a quello di un modello denso da 14 miliardi di parametri, poiché calcoliamo due moltiplicazioni di matrice da 7 miliardi di parametri.

Questa efficienza computazionale durante l’inferenza è particolarmente preziosa in scenari di distribuzione in cui le risorse sono limitate, come ad esempio dispositivi mobili o ambienti di calcolo edge. Inoltre, le ridotte esigenze computazionali durante l’addestramento possono portare a sostanziali risparmi energetici e a un minor impatto ambientale, allineandosi con l’aumento dell’attenzione per le pratiche di intelligenza artificiale sostenibili.

Sfide e considerazioni

Sebbene i modelli MoE offrano vantaggi convincenti, la loro adozione e distribuzione sono anche associate a diverse sfide e considerazioni:

1. Instabilità di addestramento: I modelli MoE sono noti per essere più soggetti a instabilità di addestramento rispetto ai loro omologhi densi. Questo problema deriva dalla natura sparsa e condizionale delle attivazioni degli esperti, che può portare a sfide nella propagazione dei gradienti e nella convergenza. Tecniche come la perdita z del router (Zoph et al., 2022) sono state proposte per mitigare queste instabilità, ma ulteriori ricerche sono ancora necessarie.
2. Finetuning e overfitting: I modelli MoE tendono a sovrapprendere più facilmente durante il finetuning, specialmente quando il compito downstream ha un set di dati relativamente piccolo. Questo comportamento è attribuito alla maggiore capacità e sparsità dei modelli MoE, che può portare a una specializzazione eccessiva sui dati di addestramento. Strategie di regolarizzazione e finetuning accurate sono richieste per mitigare questo problema.
3. Requisiti di memoria: Sebbene i modelli MoE possano ridurre i costi computazionali durante l’inferenza, spesso hanno requisiti di memoria più elevati rispetto ai modelli densi di dimensioni simili. Ciò è dovuto al fatto che tutti i pesi degli esperti devono essere caricati in memoria, anche se solo un subset è attivato per ogni input. I vincoli di memoria possono limitare la scalabilità dei modelli MoE su dispositivi con risorse limitate.
4. Bilanciamento del carico: Per raggiungere un’efficienza computazionale ottimale, è cruciale bilanciare il carico tra gli esperti, garantendo che nessun singolo esperto sia sovraccaricato mentre altri rimangono sottoutilizzati. Questo bilanciamento del carico è tipicamente raggiunto attraverso perdite ausiliarie durante l’addestramento e un’attenta regolazione del fattore di capacità, che determina il numero massimo di token che possono essere assegnati a ciascun esperto.
5. Overhead di comunicazione: Negli scenari di addestramento e inferenza distribuiti, i modelli MoE possono introdurre un overhead di comunicazione aggiuntivo a causa della necessità di scambiare informazioni di attivazione e gradienti tra gli esperti residenti su dispositivi o acceleratori diversi. Strategie di comunicazione efficienti e progettazione del modello consapevole dell’hardware sono essenziali per mitigare questo overhead.

Nonostante queste sfide, il potenziale dei modelli MoE nel consentire modelli linguistici più grandi e più capaci ha spronato notevoli sforzi di ricerca per affrontare e mitigare questi problemi.

Esempio: Mixtral 8x7B e GLaM

Per illustrare l’applicazione pratica di MoE nei modelli linguistici, consideriamo due esempi notevoli: Mixtral 8x7B e GLaM.

Mixtral 8x7B è una variante MoE del modello linguistico Mistral, sviluppato da Anthropic. Consiste in otto esperti, ciascuno con 7 miliardi di parametri, per un totale di 56 miliardi di parametri. Tuttavia, durante l’inferenza, solo due esperti sono attivati per token, riducendo efficacemente il costo computazionale a quello di un modello denso da 14 miliardi di parametri.

Mixtral 8x7B ha dimostrato prestazioni impressionanti, superando il modello Llama da 70 miliardi di parametri mentre offre tempi di inferenza molto più veloci. Una versione istruita di Mixtral 8x7B, chiamata Mixtral-8x7B-Instruct-v0.1, è stata rilasciata, migliorando ulteriormente le sue capacità nel seguire istruzioni in linguaggio naturale.

Un altro esempio degno di nota è GLaM (Google Language Model), un modello MoE di larga scala sviluppato da Google. GLaM impiega un’architettura transformer solo decoder e è stato addestrato su un enorme set di dati di 1,6 trilioni di token. Il modello raggiunge prestazioni impressionanti nelle valutazioni a pochi colpi e in un colpo solo, eguagliando la qualità di GPT-3 mentre utilizza solo un terzo dell’energia necessaria per addestrare GPT-3.

Il successo di GLaM può essere attribuito alla sua efficiente architettura MoE, che ha consentito l’addestramento di un modello con un numero vasto di parametri mantenendo esigenze computazionali ragionevoli. Il modello ha anche dimostrato il potenziale dei modelli MoE di essere più efficienti in termini energetici e più sostenibili rispetto ai loro omologhi densi.

L’architettura Grok-1

GROK MIXTURE OF EXPERT

Grok-1 è un modello MoE basato su transformer con un’architettura unica progettata per massimizzare l’efficienza e le prestazioni. Analizziamo le specifiche chiave:

1. Parametri: Con un numero impressionante di 314 miliardi di parametri, Grok-1 è il più grande modello LLM aperto fino ad oggi. Tuttavia, grazie all’architettura MoE, solo il 25% dei pesi (circa 86 miliardi di parametri) sono attivi in ogni momento, migliorando le capacità di elaborazione.
2. Architettura: Grok-1 impiega un’architettura Mixture-of-8-Experts, con ogni token elaborato da due esperti durante l’inferenza.
3. Livelli: Il modello consiste in 64 livelli di transformer, ciascuno che incorpora meccanismi di auto-attenzione e blocchi densi.
4. Tokenizzazione: Grok-1 utilizza un tokenizzatore SentencePiece con una dimensione del vocabolario di 131.072 token.
5. Incapsulamento e codifica posizionale: Il modello presenta incapsulamenti di dimensione 6.144 e utilizza incapsulamenti posizionali rotativi, consentendo un’interpretazione più dinamica dei dati rispetto alle codifiche posizionali fisse tradizionali.
6. Attenzione: Grok-1 utilizza 48 teste di attenzione per le query e 8 teste di attenzione per le chiavi e i valori, ciascuna con una dimensione di 128.
7. Lunghezza del contesto: Il modello può elaborare sequenze fino a 8.192 token di lunghezza, utilizzando la precisione bfloat16 per calcoli efficienti.

Prestazioni e dettagli di implementazione

Grok-1 ha dimostrato prestazioni impressionanti, superando LLaMa 2 70B e Mixtral 8x7B con un punteggio MMLU del 73%, mostrando la sua efficienza e accuratezza in vari test.

Tuttavia, è importante notare che Grok-1 richiede risorse GPU significative a causa delle sue dimensioni. L’attuale implementazione nel rilascio open-source si concentra sulla validazione della correttezza del modello e impiega un’implementazione del livello MoE inefficiente per evitare la necessità di kernel personalizzati.

Nonostante ciò, il modello supporta l’attivazione sharding e la quantizzazione a 8 bit, che possono ottimizzare le prestazioni e ridurre i requisiti di memoria.

In una mossa notevole, xAI ha rilasciato Grok-1 sotto la licenza Apache 2.0, rendendo i suoi pesi e l’architettura accessibili alla comunità globale per l’uso e il contributo.

Il rilascio open-source include un repository di codice di esempio JAX che dimostra come caricare e eseguire il modello Grok-1. Gli utenti possono scaricare i pesi del checkpoint utilizzando un client torrent o direttamente attraverso il HuggingFace Hub, facilitando l’accesso a questo modello innovativo.

Il futuro di Mixture-of-Experts nei modelli linguistici

Poiché la domanda di modelli linguistici più grandi e più capaci continua a crescere, l’adozione di tecniche MoE è prevista aumentare ulteriormente. Gli sforzi di ricerca in corso si concentrano sull’affrontare le sfide rimanenti, come il miglioramento della stabilità di addestramento, la mitigazione del sovrapprendimento durante il finetuning e l’ottimizzazione dei requisiti di memoria e comunicazione.

Una direzione promettente è l’esplorazione di architetture MoE gerarchiche, in cui ogni esperto è a sua volta composto da più sottoesperti. Questo approccio potrebbe potenzialmente consentire una maggiore scalabilità e efficienza computazionale mantenendo la potenza espressiva dei modelli grandi.

Inoltre, lo sviluppo di sistemi hardware e software ottimizzati per i modelli MoE è un’area di ricerca attiva. Acceleratori specializzati e framework di addestramento distribuito progettati per gestire in modo efficiente i modelli di calcolo sparsi e condizionali dei modelli MoE potrebbero ulteriormente migliorare le loro prestazioni e scalabilità.

Inoltre, l’integrazione di tecniche MoE con altri progressi nei modelli linguistici, come meccanismi di attenzione sparsi, strategie di tokenizzazione efficienti e rappresentazioni multimodali, potrebbe portare a modelli linguistici ancora più potenti e versatili in grado di affrontare una vasta gamma di compiti.

Conclusione

La tecnica Mixture-of-Experts è emersa come uno strumento potente nella ricerca di modelli linguistici più grandi e più capaci. Selezionando selettivamente gli esperti in base ai dati di input, i modelli MoE offrono una soluzione promettente alle sfide computazionali associate all’aumento delle dimensioni dei modelli densi. Sebbene ci siano ancora sfide da superare, come l’instabilità di addestramento, il sovrapprendimento e i requisiti di memoria, i benefici potenziali dei modelli MoE in termini di efficienza computazionale, scalabilità e sostenibilità ambientale li rendono un’area di ricerca e sviluppo entusiasmante.

Poiché il campo dell’elaborazione del linguaggio naturale continua a spingere i confini di ciò che è possibile, l’adozione di tecniche MoE è probabile che giochi un ruolo cruciale nell’abilitare la prossima generazione di modelli linguistici. Combinando MoE con altri progressi nell’architettura del modello, nelle tecniche di addestramento e nell’ottimizzazione dell’hardware, possiamo aspettarci modelli linguistici ancora più potenti e versatili che possano veramente comprendere e comunicare con gli esseri umani in modo naturale e fluido.