Riepilogo Tecnico: OrScale – Ottimizzazione Ortogonalizzata con Scalatura del Rapporto di Fiducia a Livello di Strato

1. Enunciato del Problema

L'ottimizzatore Muon si è affermato come uno strumento pratico per l'addestramento su larga scala, sfruttando aggiornamenti ortogonalizzati per fissare la direzione dei passi a valore matriciale tramite iterazioni di Newton–Schulz (approssimando il fattore polare). Tuttavia, Muon soffre di una limitazione critica: mentre ottimizza la direzione degli aggiornamenti, lascia la magnitudine controllata da un singolo tasso di apprendimento essenzialmente globale.

Questa scalatura globale è subottimale perché diversi strati (ad esempio, proiezioni di attenzione, proiezioni di salita/discesa MLP, embedding) esibiscono distribuzioni del gradiente e curvature efficaci molto diverse. Sebbene la variante Moonlight abbia introdotto un fattore statico per forma per equalizzare l'RMS per elemento, non può adattarsi alla deriva dinamica delle scale degli strati durante l'addestramento.

I precedenti tentativi di portare la Scalatura del Tasso di Apprendimento Adattiva a Livello di Strato (LARS) e la Scalatura del Momento Adattiva a Livello di Strato (LAMB) — che hanno disaccoppiato con successo le dimensioni dei passi dai tassi di apprendimento globali in SGD e Adam — su Muon hanno fallito. Il documento identifica che le portate naive collassano a causa di tre motivi strutturali specifici:

1. Denominatori Degeneri: L'uso della norma di Frobenius della direzione ortogonalizzata ($\|Q_\ell\|_F$) come denominatore del rapporto di fiducia fallisce perché è quasi una costante determinata dalla forma della matrice (rang), non dalla magnitudine dell'aggiornamento.
2. Disallineamento delle Unità e Saturazione del Taglio: L'uso della norma del momento grezzo ($\|\tilde{M}_\ell\|_F$) introduce un disallineamento delle unità (unità del gradiente vs unità dei pesi), causando il raggiungimento quasi universale del limite superiore di taglio del rapporto di fiducia ($r_{max}$), rendendo l'adattamento inefficace.
3. Fuga del Decadimento dei Pesi Disaccoppiato: La combinazione di denominatori calibrati con decadimento dei pesi disaccoppiato crea un ciclo di retroazione positiva in cui le norme dei pesi crescono in modo incontrollato (10–50$\times$), destabilizzando l'addestramento.

2. Metodologia

Gli autori propongono OrScale, un framework basato su un singolo principio di progettazione: Il denominatore di un rapporto di fiducia a livello di strato deve misurare la norma di Frobenius della direzione di aggiornamento reale che l'ottimizzatore sottrarrà dai parametri.

Principio Fondamentale

Sia $D_{\ell,t}$ la direzione effettiva nello spazio dei parametri sottratta dai pesi $W_{\ell,t}$ prima dell'applicazione del tasso di apprendimento globale. Il rapporto di fiducia è definito come:
$$r_{\ell,t} = \frac{\|W_{\ell,t}\|_F}{\|D_{\ell,t}\|_F + \epsilon}$$
Crucialmente, $D_{\ell,t}$ deve includere sia il termine di decadimento dei pesi sia la direzione ortogonalizzata.

Due Varianti

Il framework produce due implementazioni specifiche:

1. OrScale (Predefinito Generale/Visione):

   • Direzione di Aggiornamento: $D_{\ell,t} = \lambda W_{\ell,t} + Q_{\ell,t}$ (dove $\lambda$ è il decadimento dei pesi e $Q$ è la direzione ortogonalizzata).
   • Fattore di Forma: $s_\ell = 1$ (nessuna scalatura statica).
   • Calibrazione: Nessuna.
   • Decadimento dei Pesi: Integrato direttamente nel passo di aggiornamento scalato dal rapporto di fiducia.

2. OrScale-LM (Predefinito LLM):

   • Direzione di Aggiornamento: $D_{\ell,t} = \lambda W_{\ell,t} + s_\ell Q_{\ell,t}$.
   • Fattore di Forma: Adotta il fattore statico di Moonlight $s_\ell = 0.2\sqrt{\max(m_\ell, n_\ell)}$ per corrispondere all'RMS per elemento di AdamW.
   • Calibrazione: Introduce una costante di calibrazione per strato "lazy" ($c_{\text{denom}, \ell}$) calcolata una sola volta all'inizializzazione. Questo ancoraggio porta il rapporto di fiducia iniziale esattamente a 1 ($r_{\ell,0}=1$).
   • Significato della Calibrazione: Ciò permette a OrScale-LM di ereditare i parametri iperparametrici di tasso di apprendimento e decadimento dei pesi ottimizzati per Moonlight (e di conseguenza per AdamW) senza richiedere una nuova ricerca iperparametrica. Garantisce l'invarianza alla larghezza all'inizializzazione.

Dettagli Algoritmici

• Decadimento dei Pesi Accoppiato: Il decadimento dei pesi è integrato nel calcolo del rapporto di fiducia ($\lambda W$ fa parte del denominatore) piuttosto che applicato come un richiamo separato. Ciò previene il ciclo di retroazione di fuga osservato nelle varianti disaccoppiate.
• Efficienza: Il metodo aggiunge un sovraccarico trascurabile (<1% del tempo reale), richiedendo solo uno scalare aggiuntivo per strato matriciale e due calcoli della norma di Frobenius per passo, che sono dominati dalle esistenti iterazioni di Newton–Schulz.

3. Contributi Chiave

• Algoritmo: Una ricetta unificata e leggera sui parametri (Algoritmo 1) che integra con successo la scalatura del rapporto di fiducia a livello di strato in Muon. Gli autori identificano l'unico "angolo funzionante" nello spazio di progettazione (Denominatore della direzione di aggiornamento reale + WD accoppiato).
• Analisi dei Fallimenti: Una mappatura rigorosa dello spazio di progettazione del rapporto di fiducia di Muon (Tabella 1) che caratterizza analiticamente ed empiricamente tre modalità di fallimento (denominatore degenere, saturazione del taglio e crescita incontrollata dei pesi), dimostrando perché falliscono e perché la soluzione proposta ha successo.
• Teoria:
  • Convergenza: Dimostra una garanzia di convergenza non convessa di norma nucleare $O(1/\sqrt{T})$.
  • Guadagno Adattivo: Stabilisce una costante di discesa adattiva allo strato $\kappa_{\text{eff}} > 1$ sotto eterogeneità misurabile dello strato, migliorando strettamente la costante omogenea di Muon.
  • Separazione: Dimostra una separazione rigorosa da MuTrust/MuScale in condizioni di saturazione del taglio.
  • Calibrazione: Dimostra che OrScale-LM preserva il trasferimento del tasso di apprendimento in stile $\mu$P all'inizializzazione.
• Risultati Empirici: Valutazione completa su compiti di modelli di visione e linguaggio.

4. Risultati Sperimentali

Visione (CIFAR-10 / DavidNet)

• Prestazioni: OrScale ha raggiunto un'accuratezza top-1 di validazione di 94.05% $\pm$ 0.08, classificandosi primo tra cinque baseline.
• Miglioramento: Ciò rappresenta un miglioramento di +0.35 punti rispetto a Muon standard (93.70%) e di +0.30 rispetto a Muon+Moonlight (93.75%).
• Stabilità: OrScale ha dimostrato una finestra di tasso di apprendimento stabile più ampia rispetto a Muon, coerente con le previsioni teoriche di guadagni adattivi allo strato.
• Ablazione: Uno studio di ablazione sulle modalità di fallimento ha confermato che le varianti corrispondenti alle modalità di fallimento identificate (ad esempio, l'uso di $\|Q_\ell\|_F$ o del momento grezzo) hanno esibito i sintomi previsti (ad esempio, saturazione del taglio del 99.5–100% o fuga della norma dei pesi) e non sono riuscite a eguagliare le prestazioni di OrScale.

Grandi Modelli Linguistici (Pre-addestramento FineWeb-Edu)

• Scalabilità: Valutato su quattro scale di modello (da 125M a 1.1B parametri).
• Confronto: OrScale-LM ha superato Muon+Moonlight in tre delle quattro scale (125M, 545M, 1.1B) e ha battuto AdamW a ogni scala.
• Riduzione della Perdita: Alla scala 1.1B, OrScale-LM ha raggiunto una cross-entropy di validazione finale di 2.6251, rispetto a 2.6360 per Muon+Moonlight e 2.7304 per AdamW.
• Efficienza: Il vantaggio è stato preservato nell'intervallo di calcolo, con OrScale-LM che mostra una vittoria netta nel regime in cui il calcolo è più importante per gli LLM.

5. Significato e Affermazioni

Il documento afferma che OrScale colma con successo il divario tra la geometria spettrale di Muon e l'adattività a livello di strato di LARS/LAMB. Il suo significato principale risiede in:

1. Risoluzione del Problema della Magnitudine: Fornisce il primo controllo di magnitudine variabile nel tempo e a livello di strato per Muon che è robusto alle modalità di fallimento strutturali.
2. Trasferibilità Pratica: Attraverso la variante OrScale-LM e la calibrazione lazy, permette ai praticanti di trasferire direttamente i tassi di apprendimento ottimizzati per AdamW o Moonlight a OrScale-LM senza ri-ottimizzazione, abbassando la barriera all'adozione.
3. Solidità Teorica: Offre una giustificazione teorica rigorosa del perché i precedenti tentativi ibridi abbiano fallito e dimostra che il metodo proposto produce costanti di convergenza strettamente migliori sotto eterogeneità.
4. Validazione Empirica: Dimostra che questi miglioramenti teorici si traducono in guadagni tangibili sia in compiti di visione su piccola scala sia nel pre-addestramento di LLM su larga scala, stabilendo OrScale come un ottimizzatore all'avanguardia per strati a valore matriciale.

Gli autori mantengono modestia riguardo ai risultati sugli LLM, notando che il vantaggio alla scala 1.1B è comparabile al rumore di un singolo seme su scale più piccole, e quindi i risultati dovrebbero essere visti come evidenza direzionale dell'efficacia del metodo piuttosto che come una superiorità statistica definitiva in tutti i regimi. Tuttavia, il miglioramento coerente attraverso le scale e l'ablazione riuscita delle modalità di fallimento supportano fortemente il principio di progettazione proposto.