How Much Orthogonalization Does Muon Need?

Riepilogo Tecnico: Quanta Ortogonalizzazione Serve a Muon?

Definizione del Problema
L'ottimizzatore Muon migliora l'addestramento delle reti neurali sostituendo gli aggiornamenti del momentum mal condizionati con aggiornamenti approssimativamente semi-ortogonali. Mentre la numerica lineare standard mira a un'alta precisione nella decomposizione polare (mappando tutti i valori singolari a 1), il contesto di Muon suggerisce che un'esatta precisione possa non essere necessaria. La domanda centrale affrontata da questo lavoro è: quanta precisione di ortogonalizzazione richiede effettivamente Muon per preservare la qualità dell'addestramento? Nello specifico, il documento indaga se il costo computazionale delle iterazioni polinomiali ad alto grado (usate in metodi esistenti come Muon-Jordan e Polar Express) sia giustificato, o se un'approssimazione più economica e a basso grado possa ottenere risultati di addestramento comparabili.

Metodologia
Gli autori propongono uno schema di Newton–Schulz cubico rilassato (denominato cubic5) progettato specificamente per la banda di valori singolari a bassa precisione (bfloat16) di Muon.

1. Target Rilassati: A differenza della decomposizione polare classica che punta alla precisione di macchina, questo metodo punta a una banda rilassata di $[0.7, 1.3]$ per i valori singolari. Riconosce che in bfloat16, i valori singolari al di sotto di un limite inferiore effettivo ($l_0 = 7 \times 10^{-3}$) non sono praticamente significativi.
2. Derivazione dei Coefficienti Adattivi: Gli autori derivano un polinomio cubico adattivo $f_t(x) = a_t x + b_t x^3$ ad ogni iterazione. I coefficienti sono calcolati sulla base dell'attuale limite inferiore del valore singolare peggiore ($l_t$) e di un picco superiore fisso ($u=1.3$). Ciò assicura che il polinomio mappi l'intervallo protetto $[l_t, r_t]$ nella banda target, bilanciando al contempo gli estremi per massimizzare il nuovo limite inferiore.
3. Efficienza Computazionale:
   • Cubic5: Un singolo step cubico richiede due moltiplicazioni di matrici dominanti (formare la matrice di Gram e moltiplicarla nuovamente). Cinque step richiedono 10 moltiplicazioni di matrici dominanti.
   • Baseline (Muon-Jordan/Polar Express): Uno step quintico richiede tre moltiplicazioni di matrici dominanti. Cinque step richiedono 15 moltiplicazioni di matrici dominanti.
   • Il metodo utilizza il "trucco della trasposta" per matrici alte per minimizzare le operazioni sulla matrice di Gram più piccola.

Contributi Chiave

• Derivazione di uno Schema Cubico Rilassato: Il documento deriva uno schema di Newton–Schulz cubico a cinque step con coefficienti adattati al limite inferiore effettivo di bfloat16, raggiungendo la banda spettrale target in cinque iterazioni.
• Caratterizzazione della Frontiera Costo-Qualità: Il lavoro stabilisce che la qualità dell'addestramento non è governata monotonicamente dalla precisione della decomposizione polare. Dimostra che un'approssimazione cubica a minor costo può raggiungere la parità di prestazioni con i metodi quintici a costo maggiore e persino con la SVD esatta in specifici regimi.
• Validazione Empirica su Diverse Scale: Gli autori validano il metodo attraverso:
  • Diagnostica Sintetica: Confermare che il trasferimento del polinomio scalare regge nell'aritmetica di matrici a bassa precisione.
  • Ablazioni su NanoGPT: Confrontare cubic5, Polar Express troncato, Muon-Jordan troncato ed esplicita SVD in FP32 su GPT-2 Small.
  • Addestramento su Larga Scala: Addestramento di modelli ibridi MoE/Mamba (Nemotron-3-Nano) da 1B a 4B di parametri utilizzando Megatron-LM.

Risultati

• Benchmark Sintetici: Il metodo cubic5 mappa con successo i valori singolari nell'intervallo protetto $[7 \times 10^{-3}, 1]$ nella banda target senza esplosione numerica. I microbenchmark su GPU NVIDIA H200 mostrano che cubic5 riduce il costo della subroutine di ortogonalizzazione rispetto a Polar Express, con accelerazioni (speedup) che vanno da 1.16x a 1.52x a seconda della forma della matrice (rettangolare vs. quadrata).
• Ablazioni su GPT-2 Small:
  • Cinque step di cubic5 ottiene una perdita di validazione finale di 3.2811.
  • Cinque step di Polar Express ottiene 3.2798.
  • Cinque step di Muon-Jordan ottiene 3.2822.
  • La SVD esatta in FP32 ottiene 3.2809.
  • Osservazione: Il divario tra le varianti più forti è inferiore a 0.003 nella perdita di validazione. Notevolmente, la SVD esatta non supera i metodi polinomiali approssimati, e l'aggiunta di un quinto step a Muon-Jordan ha leggermente degradato le prestazioni rispetto a quattro step.
• Scaling Ibrido MoE/Mamba (1B–4B):
  • Sui modelli da 1B e 2B, cubic5 ha leggermente superato Muon-Jordan (perdita di validazione inferiore di $\approx 10^{-3}$).
  • Sul modello da 4B, cubic5 era leggermente superiore ($\approx 10^{-3}$), ma le traiettorie di addestramento erano quasi indistinguibili dopo la fase iniziale.
  • La differenza relativa nella perdita di addestramento è decaduta rapidamente, fluttuando intorno allo zero per il resto del budget di addestramento.

Significato e Rivendicazioni
Il documento sostiene con modestia che cubic5 è una variante pratica e a basso costo per l'ortogonalizzazione di Muon, piuttosto che una regola di aggiornamento universalmente superiore.

• Ottimizzatore vs. Precisione della Decomposizione: I risultati supportano l'ipotesi che Muon benefici di un effetto di "rimodellamento spettrale" piuttosto che di una rigorosa precisione della decomposizione polare. La mappa esatta dei valori singolari che ottimizza l'addestramento dipende dalla fase di addestramento e dal tasso di apprendimento, non solo dalla convergenza matematica a 1.
• Utilità Pratica: cubic5 offre un punto fondato sulla frontiera costo-qualità. Riduce il conteggio delle moltiplicazioni di matrici dominanti di circa un terzo (10 contro 15) mantenendo la parità della qualità empirica dell'addestramento con l'aggiornamento quintico standard di Muon-Jordan nei regimi testati.
• Limitazioni: Gli autori dichiarano esplicitamente che questi risultati stabiliscono la parità della qualità dell'addestramento, non accelerazioni (speedup) end-to-end del throughput. Lo speedup effettivo in termini di tempo di esecuzione dipende dai kernel hardware, dal batching e dalla frazione di tempo occupata dalla fase di ortogonalizzazione all'interno dell'intero ciclo dell'ottimizzatore. Il documento non afferma che cubic5 sia un algoritmo di decomposizione polare migliore di Polar Express in senso matematico, ma piuttosto un trasformatore spettrale più efficiente per il contesto specifico dell'addestramento di Muon.