I pesi strutturati rendono il training con Muon più economico su larga scala
In parole semplici
Addestrare un modello linguistico richiede risorse computazionali considerevoli. Questo paper introduce Monarch–Muon, un ottimizzatore che memorizza alcuni pesi del modello come componenti più piccole e aggiorna quelle componenti con passi matrix-aware.
In test a run singola su due tipi di modello, ha ottenuto una loss di validazione finale inferiore rispetto ad AdamW quando entrambi usavano la stessa struttura di componenti. Misurazioni separate hanno rilevato un tempo per passo di addestramento e un consumo di memoria inferiori alla dimensione più grande testata, su un singolo H100.
La ricerca
È possibile usare aggiornamenti matrix-aware nel pretraining dei modelli linguistici; aggiornamenti che possono migliorare l'ottimizzazione; senza pagare l'intero costo di storage e di calcolo dei pesi densi?
Il paper verifica se rappresentare i pesi selezionati con fattori strutturati possa rendere quegli aggiornamenti meno costosi conservandone il vantaggio di ottimizzazione.
Monarch–Muon applica aggiornamenti matrix-aware a fattori di peso strutturati, con l'obiettivo di conservare i vantaggi di ottimizzazione riducendo i costi di storage e di ortogonalizzazione.
Monarch–Muon aggiorna fattori di peso strutturati con passi matrix-aware. I ricercatori lo hanno confrontato con AdamW a strutture di fattori corrispondenti nel pretraining di modelli densi e mixture-of-experts su FineWeb-Edu.
La loss di validazione finale è risultata inferiore con Monarch–Muon in tutti e quattro i confronti a run singola. Ciò suggerisce che gli aggiornamenti matrix-aware possono essere utili all'interno di queste strutture testate.
I confronti di qualità sono run singole a dimensioni di modello specifiche, e le misurazioni separate sui sistemi non stabiliscono il compromesso combinato tra qualità e hardware in training distribuito prolungato.
La dimostrazione di convergenza assume aggiornamenti polari esatti, mentre l'ottimizzatore pratico li approssima con un numero finito di iterazioni Newton–Schulz.
Analisi
Dettagli
- Autori del paper originale
- Alexander Miasnikov, Ruslan Kabirov, Gleb Molodtsov, Dmitrii Vlashchuk, Vladimir Aletov
- Inviato ad arXiv
- Licenza
- CC BY 4.0