Muon rende più difficile bilanciare il primo layer mixture-of-experts

Abstract degli autori

Molti modelli linguistici open source basati su mixture-of-experts (MoE) vengono addestrati con Muon anziché con AdamW. La maggior parte bilancia il carico applicando un bias al punteggio di instradamento di ciascun esperto, aggiornato dopo ogni step. Le rispettive ricette non concordano sul fatto che Muon debba addestrare anche il router. In modelli con 55M e 108M di parametri attivi, misuriamo come l’ottimizzatore modifichi l’ampiezza di questi bias, cioè la differenza tra il più grande e il più piccolo in un layer. L’ampiezza aumenta quando tutti i token privilegiano gli stessi esperti, quindi indica quanto sia difficile bilanciare un layer. Con tassi di apprendimento messi a punto su esecuzioni più brevi, Muon ha aumentato l’ampiezza del primo layer MoE per ogni seed, raddoppiandola circa in media, mentre gli altri layer sono cambiati molto meno. L’aumento era dovuto all’uso di Muon nel resto della rete, indipendentemente dall’ottimizzatore usato per addestrare il router. In queste esecuzioni, gli input del primo router erano orientati nella stessa direzione e, poco dopo il warmup, l’esperto più carico riceveva in media una quota da quattro a sette volte superiore a quella che gli sarebbe spettata in caso di distribuzione equa. Nei primi 3.000 step, sottrarre la media del batch da questi input ha eliminato il sovraccarico, modificando appena la loss. Nel modello più piccolo, dimezzare il tasso di apprendimento di Muon ha eliminato gran parte dell’aumento e migliorato leggermente la loss. Usare AdamW solo per il router ha fatto variare maggiormente i carichi durante il warmup e ha aumentato leggermente la loss; a questa scala, quindi, i risultati favoriscono l’uso di Muon per il router.