Oltre due iterazioni: una ricetta scalabile per i Mixture-of-Experts con loop

Abstract degli autori

I transformer con loop introducono la profondità ricorrente come nuovo asse di scalabilità per i modelli linguistici di grandi dimensioni (LLM): applicando ripetutamente blocchi transformer condivisi, aumentano la profondità effettiva senza aumentare il numero di parametri. Tuttavia, nei grandi LLM MoE i benefici delle iterazioni restano poco chiari quando i confronti avvengono a parità di FLOP. Il motivo principale è che i vantaggi delle iterazioni aggiuntive diminuiscono rapidamente e possono persino trasformarsi in un peggioramento: i FLOP in più impiegati producono quindi pochi miglioramenti sostanziali. Di conseguenza, i lavori precedenti si fermano in genere a due iterazioni. Individuiamo due ostacoli principali alla scalabilità dei MoE con loop. In primo luogo, le iterazioni ereditano e amplificano la maledizione della profondità: la varianza degli stati nascosti cresce a ogni iterazione con l’accumularsi degli aggiornamenti residui, destabilizzando la ricorrenza profonda e causando una deriva delle rappresentazioni. In secondo luogo, i MoE con loop soffrono del collasso della selezione degli esperti: i router selezionano ripetutamente gli stessi esperti nelle diverse iterazioni, che aggiungono così calcolo senza aumentare la diversità computazionale. Sulla base di questa analisi, proponiamo LOOM, fondato su un unico principio: ogni iterazione deve apportare nuovo calcolo mantenendo stabile lo stato ricorrente. LOOM stabilizza la ricorrenza ridimensionando gli aggiornamenti residui per limitare la crescita della varianza e reiniettando l’embedding dell’input a ogni iterazione. La diversifica inoltre mediante router specifici per ciascuna iterazione, che coinvolgono esperti diversi, e un Looping Residual che porta avanti gli output delle iterazioni precedenti. Esperimenti su modelli da 100 milioni a 1,7 miliardi mostrano una scalabilità stabile fino a 9-12 iterazioni. In condizioni di quasi parità di FLOP, il modello da 700 milioni ottiene i risultati migliori con 5 iterazioni: rispetto al modello di riferimento senza loop, la perplessità scende da 18,36 a 16,54 e l’accuratezza media zero-shot sale dal 38,84% al 39,53%. Senza imporre la parità di FLOP, il modello da 1,7 miliardi addestrato su 60 miliardi di token raggiunge i risultati migliori con 9 iterazioni: la perplessità scende da 9,62 a 7,77 e l’accuratezza media zero-shot sale dal 42,4% al 47,7%. Il codice è disponibile su https://github.com/hed-ucas/LOOM.