Come iterare un MoE: appiattire gli esperti, rendere indipendenti i parametri dell’attenzione

Abstract degli autori

I transformer iterativi riutilizzano più volte un blocco di livelli: impiegando più risorse di calcolo, migliorano le prestazioni di un modello di dimensioni fisse e ne sfruttano più a fondo i parametri. I modelli sparsi mixture-of-experts (MoE), invece, attivano solo pochi esperti tra i molti disponibili per ciascun token. I MoE iterativi uniscono queste due impostazioni e offrono ai modelli MoE nuove possibilità di utilizzare meglio gli esperti, ma sollevano una domanda: come iterare un MoE? La nostra risposta è Foil. Mantenendo costanti i parametri degli esperti e il calcolo dedicato agli esperti per token, Foil (1) appiattisce gli esperti, dimezzando i livelli di esperti e raddoppiando sia gli esperti per livello sia i passaggi, così che ogni decisione di instradamento scelga da un insieme più ampio, e (2) rende indipendenti i parametri dell’attenzione, assegnandone di propri a ciascun passaggio, mentre esperti e router restano condivisi. Gli esperimenti mostrano che Foil supera nettamente il modello di riferimento iterativo senza appiattimento: a 20 miliardi di token, ogni modello Foil presenta una loss di preaddestramento inferiore; a 100 miliardi di token, la loss migliora in modo monotono all’aumentare del grado di appiattimento, e la variante di Foil più appiattita si colloca 0,012 nat sotto il modello di riferimento a parità di parametri e di calcolo, con un’accuratezza nei compiti a valle pari o superiore. Rendere indipendenti i parametri dell’attenzione produce inoltre un instradamento più equilibrato e più sicuro a parità di struttura. Le nostre analisi delle singole componenti chiariscono perché Foil funziona e traducono i risultati in indicazioni per progettare MoE iterativi: i vantaggi dell’iterazione e dell’aumento degli esperti per livello si rafforzano a vicenda; la sicurezza delle decisioni di instradamento rispecchia un buon utilizzo degli esperti meglio dell’equilibrio del carico; un MoE sparso iterativo dovrebbe quindi usare più esperti per livello e più passaggi. Codice e configurazioni sono disponibili all’indirizzo https://github.com/SR-A-W/how-to-loop-moe.