Addestramento dei mixture-of-experts robusto rispetto alla distribuzione

Abstract degli autori

I transformer Mixture-of-Experts (MoE) aumentano la propria capacità attivando solo pochi esperti per token, ma questa attivazione selettiva crea un problema di affidabilità nascosto: quando l’instradamento è imperfetto, i modelli con carico bilanciato possono inviare i token a esperti non sufficientemente addestrati per gli input loro assegnati. Proponiamo Distributionally Robust MoE Training (DRMoET), una funzione obiettivo integrabile senza modifiche all’architettura che considera gli esperti di ciascun livello come gruppi di robustezza definiti dal modello stesso e ottimizza gli esiti dell’instradamento con perdite elevate, anziché limitarsi a uniformare il traffico. DRMoET aggiorna la distribuzione degli esperti per ciascun livello mediante una regola softmax regolarizzata con l’entropia, applicata alle perdite degli esperti ponderate in base alle attivazioni e smussate tramite media mobile esponenziale (EMA). In questo modo rafforza i percorsi plausibili di instradamento verso esperti non al primo posto, preservando il calcolo standard dei MoE. Seguendo la procedura FLAME-MoE su scale di 746 milioni e 10,3 miliardi di parametri totali, DRMoET migliora le medie nei compiti a valle rispetto sia a FLAME-MoE standard sia al bilanciamento senza perdita ausiliaria. Con 10,3 miliardi di parametri totali e 67 miliardi di token di addestramento, DRMoET porta la media su sette compiti da 0,6625 a 0,6767, mentre il modello di riferimento senza perdita ausiliaria raggiunge 0,6431. Le analisi dei meccanismi mostrano una minore varianza delle perdite degli esperti, con una perdita media quasi invariata, una perdita in eccesso inferiore del 4,3% in caso di instradamento errato forzato verso posizioni intermedie di $k$ e una migliore specializzazione degli esperti per dominio. Questi risultati indicano nella robustezza dell’instradamento, e non solo nel bilanciamento dell’utilizzo, un obiettivo pratico per aumentare in modo affidabile la scala dei MoE sparsi. La pagina del progetto e il codice sono disponibili all’indirizzo: https://drmoet.github.io/.