Muon+: Towards Better Muon via One Additional Normalization Step

Il lavoro propone Muon+, un'ottimizzazione semplice ma efficace dell'ottimizzatore Muon che introduce un passo di normalizzazione aggiuntivo dopo l'ortogonalizzazione, dimostrando attraverso estesi esperimenti di pre-addestramento su modelli di grandi dimensioni un miglioramento coerente delle prestazioni rispetto al Muon originale.

L'essenziale

🚀 Il Problema: Costruire un'Intelligenza Artificiale Gigante

Immagina di dover costruire un grattacielo altissimo (una Intelligenza Artificiale o LLM) usando mattoni digitali. Più il palazzo è alto, più è difficile assicurarsi che non crolli mentre lo costruisci.

Per costruire questi "palazzi digitali", gli scienziati usano degli ottimizzatori. Pensali come gli ingegneri di cantiere che decidono come spostare i mattoni per rendere la struttura più solida.
Fino a poco tempo fa, il capo cantiere più famoso era un certo Adam (o AdamW). Ma recentemente è arrivato un nuovo ingegnere molto intelligente chiamato Muon.

Cos'ha di speciale Muon?
Muon è come un ingegnere che ha un superpotere: sa raddrizzare i mattoni. Quando i mattoni (i dati) sono tutti storti o ammassati in modo disordinato, Muon li allinea perfettamente l'uno rispetto all'altro (un processo chiamato "ortogonalizzazione"). Questo aiuta l'edificio a crescere più velocemente e senza crolli.

💡 La Scoperta: Il "Tocco in Più" di MUON+

Gli autori di questo studio (Ruijie Zhang e il suo team) si sono chiesti: "Muon è già bravo, ma possiamo renderlo ancora meglio?".

Hanno notato che dopo aver allineato i mattoni, Muon li lasciava così, senza controllarne la grandezza. Immagina di avere una squadra di operai che allineano perfettamente i mattoni, ma alcuni sono enormi e altri minuscoli. Il risultato? La struttura è dritta, ma instabile.

La loro idea geniale è stata aggiungere un solo passaggio in più: Normalizzare.
In parole povere, dopo aver allineato i mattoni, MUON+ dice: "Ehi, fermiamoci un attimo e assicuriamoci che tutti i mattoni abbiano la stessa dimensione perfetta prima di metterli in posizione".

È come se, dopo aver sistemato le ruote di un'auto (Muon), un meccanico aggiuntivo (MUON+) controllasse che la pressione degli pneumatici fosse perfetta su tutte e quattro le ruote prima di partire.

🏗️ Come l'hanno testato?

Gli scienziati hanno messo alla prova questo nuovo metodo "MUON+" costruendo diversi tipi di edifici digitali:

1. Piccoli cottage (modelli da 130 milioni di parametri).
2. Grattacieli medi (modelli da 1 miliardo di parametri).
3. Città intere (modelli molto grandi).

Hanno usato due stili di architettura famosi: GPT (come ChatGPT) e LLaMA (un altro modello molto popolare).

📈 I Risultati: Perché è meglio?

I risultati sono stati sorprendenti e costanti:

• Più veloce e stabile: Con MUON+, i modelli imparano meglio e fanno meno errori durante l'addestramento.
• Funziona ovunque: Che tu stia costruendo un piccolo cottage o un grattacielo, MUON+ funziona sempre meglio del vecchio Muon.
• Resiste alla fatica: Hanno anche provato a far lavorare i modelli per tempi lunghissimi (addestramento "overtraining", come se dovessero costruire un intero quartiere invece di una sola casa). Anche in questi casi estremi, MUON+ non si è stancato e ha mantenuto la sua efficienza.

🔍 L'Analogia Finale: La Squadra di Calciatori

Immagina di allenare una squadra di calcio (il modello AI).

• Muon è l'allenatore che insegna ai giocatori a passare la palla in modo coordinato, assicurandosi che non si urtino a vicenda (allineamento/ortogonalizzazione).
• MUON+ è lo stesso allenatore, ma che aggiunge un dettaglio fondamentale: dopo aver insegnato le manovre, controlla che tutti i giocatori abbiano la stessa energia e la stessa forza prima di scendere in campo.

Senza questo controllo (normalizzazione), alcuni giocatori potrebbero correre troppo e stancarsi, mentre altri potrebbero essere troppo lenti. Con MUON+, la squadra è equilibrata, coordinata e pronta a vincere.

🏁 Conclusione

In sintesi, questo paper ci dice che per costruire le Intelligenze Artificiali del futuro, non serve sempre inventare macchine complesse e costose. A volte, basta un piccolo aggiustamento intelligente (come quel passaggio di normalizzazione in più) per ottenere risultati molto più grandi, risparmiando tempo e risorse.

È come scoprire che per fare il caffè perfetto, non serve una macchina nuova, ma basta aggiungere un secondo di attesa dopo aver schiacciato il caffè: il risultato è semplicemente migliore.

Sommario tecnico

1. Il Problema

L'addestramento di modelli fondazionali su larga scala (come GPT, LLaMA, DeepSeek) sta diventando proibitivamente costoso in termini computazionali. Sebbene ottimizzatori come Adam e AdamW siano dominanti, la ricerca si sta concentrando su metodi più efficienti per ridurre i costi di calcolo e memoria.

L'ottimizzatore Muon ha recentemente dimostrato prestazioni promettenti nel pre-addestramento di grandi modelli linguistici. La sua idea chiave è l'ortogonalizzazione della matrice del momento (gradient momentum) tramite iterazioni di Newton-Schulz, una tecnica progettata per contrastare il collasso del rango dei gradienti. Tuttavia, nonostante i successi di Muon, gli autori notano che esiste ancora spazio per migliorarne la stabilità e l'efficienza, specialmente in regimi di addestramento computazionalmente ottimali o con rapporti token-parametro (T2P) molto elevati.

2. Metodologia: MUON+

Gli autori propongono MUON+, un'evoluzione semplice ma efficace di Muon. La modifica fondamentale consiste nell'introdurre un passo di normalizzazione aggiuntivo immediatamente dopo l'ortogonalizzazione della matrice del momento, ma prima dell'aggiornamento dei pesi.

L'Algoritmo

Mentre Muon standard aggiorna i pesi $W_t$ come:
$$W_t = W_{t-1} - \eta \cdot \sqrt{m/n} \cdot \text{Ortho}(M_t)$$

MUON+ introduce l'operatore di normalizzazione $\text{Norm}^{(d)}$ lungo una direzione $d$:
$$W_t = W_{t-1} - \eta \cdot \sqrt{m/n} \cdot \text{Norm}^{(d)}(\text{Ortho}(M_t))$$

L'operatore $\text{Ortho}(\cdot)$ approssima la parte polare della matrice (tramite iterazioni di Newton-Schulz) per ottenere una matrice semi-ortogonale. Successivamente, viene applicata una normalizzazione $\ell_2$ lungo una specifica direzione:

• Normalizzazione per colonna (col): Normalizza ogni colonna della matrice.
• Normalizzazione per riga (row): Normalizza ogni riga della matrice.
• Combinazioni: È possibile applicare sequenzialmente normalizzazioni per riga e colonna (es. col_row o row_col).

Gli autori ipotizzano che la normalizzazione strutturale degli aggiornamenti ortogonali sia il vero motore dei miglioramenti di prestazioni, più che componenti aggiuntive complesse come l'adattamento del secondo momento o formulazioni su varietà (manifold).

3. Contributi Chiave

1. Semplicità ed Efficacia: MUON+ aggiunge un solo passo computazionale (normalizzazione) all'algoritmo Muon esistente, mantenendo la complessità quasi invariata.
2. Validazione su Scala Industriale: Il metodo è stato testato su un'ampia gamma di architetture (GPT-style e LLaMA-style) e dimensioni, dai 60M ai 1,3 miliardi di parametri.
3. Regimi di Addestramento Estesi: Gli esperimenti coprono non solo il regime "compute-optimal" (rapporto T2P $\approx 20$), ma estendono il rapporto fino a livelli industriali di $\approx 200$ (overtraining), dimostrando la scalabilità del metodo.
4. Analisi Ablativa: Gli autori hanno isolato il contributo della normalizzazione rispetto ad altre modifiche (come quelle presenti in NorMuon o Mano), dimostrando che la normalizzazione è il fattore dominante per i guadagni di prestazioni.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti sul dataset FineWeb utilizzando GPU H100/A100 in precisione mista (bfloat16).

• Modelli GPT (130M - 774M parametri):

  • MUON+ ha superato costantemente Muon in termini di Perplexity di validazione (PPL).
  • Riduzioni della PPL di 2.02 per GPT-Small e 1.72 per GPT-Base.
  • Anche per GPT-Large (774M) si è osservato un miglioramento significativo (-0.91 PPL).

• Modelli LLaMA (60M - 1B parametri):

  • MUON+ ha ottenuto le migliori prestazioni complessive rispetto sia ad AdamW che a Muon standard.
  • Per LLaMA-1B, la PPL è scesa da 10.68 (Muon) a 10.31 (MUON+).

• Overtraining (Rapporto T2P $\approx 200$):

  • In esperimenti di overtraining su 72 miliardi di token, MUON+ ha mantenuto un vantaggio stabile rispetto a Muon, con riduzioni della PPL di 1.13 per GPT-Base e 0.45 per LLaMA-350M.
  • Le curve di perdita mostrano che MUON+ scala favorevolmente con l'aumento dei token di addestramento senza degradare nelle fasi finali.

• Robustezza e Iperparametri:

  • Direzione di normalizzazione: Le combinazioni col_row e row_col hanno generalmente fornito le migliori prestazioni. La normalizzazione per riga (row) si è rivelata leggermente superiore a quella per colonna (col).
  • Learning Rate: MUON+ mostra una minore sensibilità alla scelta del learning rate, mantenendo prestazioni stabili anche con learning rate subottimali (troppo alti), a differenza di Muon standard.
  • Indipendenza dal metodo di ortogonalizzazione: I benefici di MUON+ sono stati confermati utilizzando diversi metodi di approssimazione SVD (You, Jordan, PolarExpress), dimostrando che il miglioramento è agnostico rispetto al metodo di ortogonalizzazione specifico.

5. Significato e Conclusione

Il paper MUON+ stabilisce che la normalizzazione strutturale degli aggiornamenti ortogonali è un componente critico, spesso sottovalutato, per la stabilità e l'efficienza dell'addestramento di grandi modelli linguistici.

• Impatto Pratico: Offre un miglioramento "drop-in" (facilmente integrabile) che richiede solo un passo aggiuntivo, portando a modelli finali di qualità superiore e a una convergenza più stabile.
• Scalabilità: Dimostra che i benefici non sono limitati a piccoli modelli o regimi di calcolo ottimali, ma persistono e si rafforzano in scenari di addestramento massiccio (overtraining) tipici dell'industria.
• Implicazioni Teoriche: Suggerisce che la chiave per migliorare ottimizzatori basati su ortogonalizzazione (come Muon) risiede nella gestione della norma degli aggiornamenti, piuttosto che in meccanismi complessi di adattamento del secondo momento o proiezioni su varietà.

In sintesi, MUON+ rappresenta un passo avanti significativo nell'ottimizzazione per l'addestramento di foundation models, offrendo un metodo semplice, robusto e scalabile che supera lo stato dell'arte attuale (Muon) su una vasta gamma di configurazioni.