Cos'è il mixture of experts?

Le architetture Mixture of Experts consentono ai modelli su larga scala, anche quelli che comprendono diversi miliardi di parametri, di ridurre notevolmente i costi di elaborazione durante il pre-addestramento e di ottenere prestazioni più rapide durante il tempo di inferenza. In generale, raggiunge questa efficienza attivando selettivamente solo gli esperti specifici necessari per un determinato compito, anziché attivare l'intera rete neurale per ogni attività.

Sebbene gran parte delle implementazioni moderne delle configurazioni mixture of experts sia stata sviluppata (all'incirca) nell'ultimo decennio, la premessa fondamentale dei modelli MoE ha origine nell'articolo del 1991 "Adaptive Mixture of Local Experts". Il documento proponeva di addestrare un sistema AI composto da reti separate, ognuna specializzata in un diverso sottoinsieme di casi di addestramento. Ciò comportava l'addestramento sia delle "reti di esperti" stesse che di una rete di gating che determinava quale esperto doveva essere utilizzato per ogni attività secondaria. Rispetto a un modello convenzionale analogo, gli autori hanno riscontrato che il loro modello sperimentale era significativamente più veloce da addestrare: raggiungeva la soglia di accuratezza desiderata in un numero di epoch di addestramento pari alla metà di quello del modello convenzionale.¹

Negli ultimi anni, poiché i principali modelli di deep learning utilizzati per l'AI generativa sono diventati sempre più grandi e impegnativi dal punto di vista dell'elaborazione, il mixture of experts offre un mezzo per affrontare il compromesso tra la maggiore capacità dei modelli più grandi e la maggiore efficienza dei modelli più piccoli. Questo aspetto è stato esplorato in particolare nel campo dell'elaborazione del linguaggio naturale (NLP): alcuni dei principali modelli linguistici di grandi dimensioni (LLM), come Mixtral 8x7B di Mistral e (secondo alcuni report) GPT-4² di OpenAI, hanno utilizzato l'architettura MoE.

I moderni modelli di deep learning sono costruiti a partire da reti neurali artificiali, che comprendono più livelli di nodi interconnessi (o "neuroni"). Ogni neurone ha una funzione di attivazione: un'operazione matematica eseguita sui dati ricevuti dal livello precedente, il cui output alimenta l'input inviato al livello successivo. Le reti neurali feed-forward (FFN) classiche elaborano le informazioni passando progressivamente i dati di input dai neuroni di un livello ai neuroni del livello successivo, fino a raggiungere un livello esterno dove si verificano le previsioni finali. Alcune architetture di reti neurali incorporano elementi aggiuntivi, come i meccanismi di auto-attenzione dei modellidi trasformazione , che acquisiscono pattern e dipendenze aggiuntivi nei dati di input.

Le connessioni tra livelli e neuroni diversi sono mediate da parametri del modello che possono essere appresi: pesi (e bias) variabili che amplificano o diminuiscono l'influenza che una determinata parte dell'output della rete ha su altre parti della rete. Un modello di deep learning "impara" regolando questi parametri, utilizzando algoritmi di ottimizzazione come la discesa del gradiente, in modo da aumentare l'accuratezza delle sue previsioni.

Se da un lato un numero maggiore di parametri aumenta la capacità del modello, ovvero la sua capacità di assorbire le informazioni e i pattern in esso contenuti, dall'altro aumenta anche le risorse computazionali necessarie per addestrare e far funzionare il modello. In un tipico modello di deep learning, che in questo contesto viene definito modello denso, viene eseguita l'intera rete per elaborare tutti gli input. In questo modo, si crea un compromesso tra la capacità e la praticità del modello.

A differenza dei modelli densi convenzionali, il mixture of experts utilizza la conditional computation per imporre la scarsità: anziché utilizzare l'intera rete per ogni input, i modelli MoE apprendono una funzione di mappatura computazionalmente poco costosa che determina quali porzioni della rete (in altre parole, quali esperti) sono più efficaci per elaborare un determinato input, come un singolo token utilizzato per rappresentare una parola o un frammento di parola nelle attività di NLP.

Ciò consente di aumentare la capacità del modello (espandendo il numero totale di parametri) senza un corrispondente aumento del carico di elaborazione richiesto per addestrarlo ed eseguirlo (perché non tutti questi parametri saranno necessariamente utilizzati in un dato momento).

I modelli MoE elaborano i dati designando una serie di "esperti", ognuno dei quali è una sottorete all'interno di una rete neurale più grande, e addestrando una rete di gating (o router) ad attivare solo gli esperti specifici più adatti a un determinato input.

Il vantaggio principale dell'approccio MoE è che applicando la scarsità, anziché attivare l'intera rete neurale per ogni token di input, è possibile aumentare la capacità del modello mantenendo essenzialmente costanti i costi computazionali.

A livello di architettura, ciò si ottiene sostituendo i tradizionali livelli di rete feed-forward (FFN) densi con livelli MoE sparsi (o blocchi). Nel gergo delle reti neurali, "blocco" si riferisce a un elemento strutturale ricorrente che svolge una funzione specifica. In un modello MoE sparso (SMoE), questi blocchi di esperti possono essere livelli singoli, FFN autonomi o persino MoE annidati su se stessi.

Ad esempio, nel modello linguistico Mixtral 8x7B di Mistral, ogni livello è composto da 8 blocchi feedforward, cioè esperti, ognuno dei quali ha 7 miliardi di parametri. Per ogni token, a ogni livello, una rete di router seleziona due di questi otto esperti per elaborare i dati. Combina quindi gli output di questi due esperti e trasmette il risultato al livello successivo. Gli esperti specifici selezionati dal router in un determinato livello possono essere diversi da quelli selezionati al livello precedente o successivo.³

Uno SMoE può essere composto interamente da livelli MoE sparsi, ma molte architetture di modelli MoE includono blocchi sia sparsi che densi. Questo vale anche per Mixtral, in cui i blocchi responsabili del meccanismo di auto-attenzione del modello sono condivisi da tutti gli 8 esperti. In pratica, questo rende denominazioni come "8x7B" potenzialmente fuorvianti: poiché molti dei parametri del modello sono condivisi da ciascuna sottorete di esperti da 7 miliardi di parametri, Mixtral ha un totale di circa 47 miliardi di parametri, non 56 miliardi come si potrebbe supporre con una semplice moltiplicazione.

Questo conteggio complessivo dei parametri è comunemente indicato come conteggio dei parametri sparsi e può essere generalmente inteso come una misura della capacità del modello. Il numero di parametri che saranno effettivamente utilizzati per elaborare un singolo token (mentre transita attraverso alcuni blocchi di esperti e ne aggira altri) si chiama conteggio dei parametri attivi e può essere inteso come una misura dei costi computazionali del modello. Sebbene ciascun token immesso in Mixtral abbia accesso a 46,7 miliardi di parametri, per elaborare un determinato esempio vengono utilizzati solo 12,9 miliardi di parametri attivi.

La comprensione di questo utilizzo ottimale dei conteggi dei parametri è fondamentale per capire i vantaggi dei modelli MoE. Ad esempio, Mixtral supera la variante da 70 miliardi di parametri di Llama 2 di Meta nella maggior parte dei benchmark, con una velocità decisamente maggiore, nonostante abbia un numero totale di parametri inferiore di un terzo e utilizzi meno del 20% del numero di parametri attivi al momento dell'inferenza.³

Vale la pena notare, tuttavia, che il conteggio complessivo dei parametri di un MoE sparso non è del tutto irrilevante per i requisiti computazionali. Nonostante l'utilizzo di un solo sottoinsieme di parametri durante l'inferenza, la maggior parte dei parametri del modello viene caricata in memoria, il che significa che l'efficienza computazionale di cui godono i MoE nella maggior parte dei casi non si applica ai loro requisiti di RAM/VRAM.

La chiave del concetto (e dell'efficienza) dei MoE è che solo alcuni degli esperti (e quindi dei parametri) di un livello sparso saranno attivati in un dato momento, riducendo così i requisiti di elaborazione attivi.

Sebbene la conditional computation sia stata a lungo proposta come mezzo teorico per disaccoppiare le richieste computazionali dall'aumento della capacità del modello, le sfide algoritmiche e prestazionali per la sua corretta esecuzione non sono state superate fino all'articolo di Shazeer et al del 2017 "Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer".⁴

I vantaggi dei livelli sparsi rispetto a quelli densi sono più evidenti quando si ha a che fare con dati ad alta dimensionalità, in cui i modelli e le dipendenze sono spesso complessi e non lineari: per esempio, nelle attività di NLP che richiedono a un modello di elaborare una lunga sequenza di testo, ogni parola è tipicamente correlata solo a un piccolo sottoinsieme di altre parole in quella sequenza. Questo rende gli SMoE un'area di enorme potenziale nel campo degli LLM, in cui modelli MoE ben calibrati possono godere dei vantaggi della scarsità senza sacrificare le prestazioni. I modelli MoE "sparsely-gated" sono stati applicati con successo anche alle attività di computer vision,^{5 6} e rimangono un'area di studio attiva in questo campo. 

Questa scarsità si ottiene attraverso la conditional computation: l'attivazione dinamica di parametri specifici in risposta a input specifici. Una progettazione efficace della rete di gating (o "router"), che impone la conditional computation, è quindi essenziale per il successo dei modelli MoE.

È possibile utilizzare numerosi meccanismi di controllo per selezionare quali esperti utilizzare in una determinata situazione. La funzione di gating è fondamentale per le prestazioni del modello, in quanto una strategia di routing inadeguata può far sì che alcuni esperti siano poco addestrati o eccessivamente specializzati, riducendo l'efficacia dell'intera rete.

Un tipico meccanismo di gating in una configurazione MoE tradizionale, introdotto nel famoso articolo di Shazeer, utilizza la funzione softmax: per ciascuno degli esperti, per ogni singolo esempio, il router prevede un valore di probabilità (basato sui pesi delle connessioni di quell'esperto al parametro corrente) di quell'esperto che produce il miglior output per un dato input: anziché calcolare l'output di tutti gli esperti, il router calcola solo l'output di (quelli che prevede essere) i migliori k esperti per quell'esempio. Come già descritto, Mixtral utilizza questa classica strategia di routing top-k: in particolare, utilizza il routing top-2, ovvero k=2, selezionando i 2 migliori (su un totale di 8) esperti.

Nel loro famoso articolo del 2021, "Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficiency Sparsity", Fedus et al. hanno portato il routing top-k all'estremo: lavorando con l'LLM T5 di Google, hanno sostituito i livelli FFN del modello con 128 esperti e hanno implementato k=1, chiamato anche "hard routing". Anche portando il modello fino a un trilione di parametri, questa configurazione ha migliorato la velocità di pre-addestramento del 400%.⁶

Nonostante i loro numerosi vantaggi, i MoE rendono il processo di addestramento notevolmente più complesso. Un importante aspetto negativo della strategia di routing "vanilla" top-k è la possibilità che la rete di gating converga per attivare solo pochi esperti. Questo è un problema che si autoalimenta: se una manciata di esperti viene selezionata precocemente in modo sproporzionato, quegli esperti verranno addestrati più rapidamente e poi continueranno a essere selezionati in quanto producono previsioni più affidabili rispetto agli altri esperti meno addestrati. Questo carico sbilanciato fa sì che gli altri esperti finiscano, in senso figurato e letterale, per essere dei pesi morti.

Per mitigare questo problema, Shazeer et al. hanno introdotto il gating top-k rumoroso: si aggiunge un po' di rumore gaussiano ai valori di probabilità previsti per ciascun esperto, introducendo una certa casualità che incoraggia un'attivazione più uniformemente distribuita degli esperti. Hanno anche aggiunto due termini di regolarizzazione addestrabili alla selezione degli esperti: ridurre al minimo la perdita di bilanciamento del carico penalizza l'eccessivo affidamento su un singolo esperto, mentre ridurre al minimo la perdita di diversità degli esperti premia un utilizzo equo di tutti gli esperti.

Un paper di Google del 2020, "GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding", ha introdotto due ulteriori metodi di bilanciamento del carico:

• Random routing: mentre l'esperto "migliore" nella configurazione top-2 viene selezionato utilizzando la funzione softmax standard, il secondo esperto viene scelto in modo semi-casuale (con la probabilità che un esperto venga scelto in proporzione al peso della sua connessione). Il secondo esperto con il punteggio più alto ha quindi maggiori probabilità di essere selezionato, ma non è più garantito che venga selezionato.

• Expert capacity: gli autori stabiliscono una soglia che definisce il numero massimo di token che possono essere elaborati da un singolo esperto. Se uno dei primi due esperti scelti ha raggiungo la capacità massima, il token viene considerato "in overflow" e passa al livello successivo della rete.⁷

Come accennato in precedenza, i vantaggi dei modelli sparsi sono talvolta attenuati dalla loro maggiore complessità. Le sfide legate all'implementazione dei MoE sono particolarmente evidenti nel processo di tuning (messa a punto). I modelli sparsi sono più inclini all'overfitting rispetto ai modelli densi tradizionali e la presenza sia di livelli MoE sparsi che di livelli FFN densi complica un approccio unico per tutti.

Sono state proposte diverse osservazioni e approcci per ridurre l'instabilità durante la messa a punto dei MOE. Gli autori dell'articolo di Switch Transformers hanno osservato che le varianti con meno esperti hanno goduto di una messa a punto più efficace, il che suggerisce che i vantaggi di un numero maggiore di esperti nel pre-addestramento possono essere contrastati dal suo ostacolo alla specializzazione nelle attività a valle. 

Nell'articolo del 2022 "ST-MoE: Design Stable and Transferable Sparse Expert Models", Zoph et al hanno confrontato i risultati di 5 diversi approcci: messa a punto di tutti i parametri ("All"), dei soli parametri non MoE ("Non MoE"), dei soli parametri MoE ("MoE"), dei soli parametri di auto-attenzione e dei parametri di attenzione encoder-decoder ("Attention") e dei soli parametri FFN non MoE ("FFN").

• Non è stata riscontrata quasi alcuna differenza tra All e Non-MoE

• La messa a punto dei soli parametri di attenzione ha comportato una leggera riduzione delle prestazioni.

• L'aggiornamento dei soli parametri MoE ha ridotto significativamente le prestazioni del modello, nonostante circa l'80% dei parametri del modello risiedesse nei livelli MoE sparsi.

• FFN è stato l'unico approccio che ha migliorato le prestazioni rispetto alla baseline All. 

Gli autori hanno ipotizzato che, poiché i livelli esperti rappresentavano solo un quarto dei livelli totali del loro modello e un token vedrà al massimo solo due esperti per livello, l'isolamento dei parametri MoE comporta aggiornamenti del peso meno completi (e quindi maggiore overfitting e perdita di addestramento).⁸

Un articolo di luglio 2023, "Mixture-of-Experts Meets Instruction Tuning", ha esplorato l'impatto dell'instruction tuning sui modelli MoE utilizzando equivalenti di Google T5 e Flan-T5, una versione di T5 sottoposta a instruction tuning con il protocollo Flan di Google, come baseline. Il loro esperimento ha confrontato quattro configurazioni: messa a punto di un modello T5 denso, messa a punto di un modello Flan-T5 denso, messa a punto di un modello MoE e messa a punto di un modello Flan-MoE sottoposto a instruction tuning.

Come previsto, l'equivalente di T5 denso ha superato il MoE dopo la messa a punto. Tuttavia, il modello Flan-MoE messo a punto ha notevolmente superato il modello Flan-T5 messo a punto. Inoltre, il miglioramento di Flan-MoE rispetto al MoE è stato persino maggiore del miglioramento di Flan-T5 rispetto al T5 originale.9

Questo suggerisce in modo incoraggiante che, nonostante le loro difficoltà con la messa a punto standard, i modelli MoE in realtà beneficiano maggiormente dell'instruction tuning rispetto alle loro controparti dense. Questa scoperta è stata realizzata con l'acclamato rilascio di Mixtral 8x7B Instruct, una variante di Mixtral sottoposta a instruction tuning offerta come foundation model in IBM watsonx.ai.