Riepilogo Tecnico: Muown - Controllo della Norma Righe per l'Ottimizzazione Muon

Enunciato del Problema

Muon si è affermato come un ottimizzatore competitivo per il pre-addestramento di modelli linguistici su larga scala (LLM), spesso superando AdamW sfruttando la struttura matriciale dei livelli lineari dei Transformer. Tuttavia, osservazioni recenti indicano che il comportamento di Muon su larga scala è sensibile al decadimento dei pesi (weight decay) e soggetto a instabilità numeriche, come overflow in bfloat16 e logit dell'attenzione esplosivi.

Un problema critico identificato è la costante deriva verso l'alto della norma spettrale ($\|W_t\|_{S_\infty}$) delle matrici dei pesi durante l'addestramento quando Muon viene utilizzato senza decadimento dei pesi disaccoppiato. Questa deriva viola la condizione spettrale richiesta per un apprendimento delle caratteristiche stabile ($\|W\|_{S_\infty} = \Theta(\sqrt{m/n})$). Sebbene il decadimento dei pesi disaccoppiato possa mitigare questo fenomeno riducendo uniformemente la matrice dei pesi, si tratta di un intervento non mirato che scala l'intero spettro indipendentemente dalle direzioni specifiche che guidano la crescita, potenzialmente rallentando la convergenza nelle fasi iniziali dell'addestramento. Il lavoro mira a identificare il driver empirico specifico di questa deriva della norma spettrale e a proporre un meccanismo di controllo più mirato.

Metodologia: Muown

Gli autori introducono Muown, una sostituzione diretta (drop-in replacement) per Muon che affronta la deriva della norma spettrale attraverso una strategia di riparametrizzazione anziché vincoli aggregati.

1. Decomposizione della Norma Spettrale

L'idea centrale è una decomposizione della norma spettrale in due fattori (Proposizione 1):
$$\|W_t\|_{S_\infty}^2 = \|g_t\|_\infty^2 \cdot \lambda_{\max}(P_t C_t P_t)$$
Dove:

• $g_t = \|W_t\|_{row}$ è il vettore delle magnitudini delle righe (norme euclidee delle righe).
• $D_t$ è la matrice normalizzata per riga ($W_t = \text{Diag}(g_t)D_t$).
• $C_t = D_t D_t^\top$ è la matrice di Gram delle righe a norma unitaria.
• $P_t$ è una matrice diagonale delle magnitudini normalizzate delle righe.
• $\lambda_{\max}(P_t C_t P_t)$ rappresenta il fattore di coerenza delle righe.

L'analisi empirica su modelli da 500M rivela che il fattore di magnitudine delle righe ($\|g_t\|_\infty$) è il driver principale della deriva della norma spettrale, muovendosi in sincronia con $\|W_t\|_{S_\infty}$. Al contrario, il fattore di coerenza delle righe rimane ben comportato e limitato durante tutto l'addestramento.

2. Riparametrizzazione e Ottimizzazione

Basandosi su questa diagnosi, Muown riparametrizza la matrice dei pesi $W$ come:
$$W(g, R) = \text{Diag}\left(\frac{g}{\|R\|_{row}}\right) R$$
Qui, $g \in \mathbb{R}^m$ rappresenta il vettore esplicito delle magnitudini delle righe, e $R \in \mathbb{R}^{m \times n}$ rappresenta il componente direzionale.

• Componente Direzionale ($R$): Muon viene applicato invariato a $R$, trattandolo sotto la geometria della norma spettrale (discesa spettrale più ripida).
• Componente di Magnitudine ($g$): Le magnitudini delle righe $g$ sono promosse a variabili esplicite dell'ottimizzatore. Vengono aggiornate sotto la geometria $\ell_\infty$, che è la norma duale naturale per la norma dell'operatore delle matrici diagonali.
• Regola di Aggiornamento: L'aggiornamento per $g$ utilizza un ottimizzatore Adam (che agisce come SignSGD smussato sotto la geometria $\ell_\infty$) per gestire il rumore stocastico, mentre l'aggiornamento per $R$ segue la procedura standard di Muon con ortogonalizzazione Newton-Schulz.

Questo approccio è implementato implicitamente all'interno dello stato dell'ottimizzatore, lasciando invariato il passaggio in avanti (forward pass) dell'architettura del modello.

Contributi Chiave

1. Diagnosi della Deriva Spettrale: Il lavoro decompone la norma spettrale per identificare la massima magnitudine della riga ($\|g_t\|_\infty$) come driver empirico della crescita della norma spettrale sotto Muon, distinguendola dal fattore di coerenza delle righe stabile.
2. Algoritmo Muown: Un nuovo ottimizzatore che tratta le magnitudini delle righe come variabili addestrabili esplicite aggiornate sotto la geometria $\ell_\infty$, preservando al contempo gli aggiornamenti direzionali di Muon. Questo evita la riduzione indiscriminata del decadimento dei pesi e le prior di ortogonalità delle riparametrizzazioni spettrali.
3. Garanzie di Convergenza: Gli autori dimostrano che Muown ottiene tassi di convergenza ottimali non convessi sia in regimi deterministici che stocastici. L'analisi è condotta in una norma duale mista ($\ell_1$ per le magnitudini, Schatten-1 per le direzioni) allineata con le geometrie di aggiornamento. Il limite teorico suggerisce una riduzione del coefficiente di rumore stocastico efficace rispetto a Muon.
4. Validazione Empirica: Esperimenti estensivi su FineWeb-Edu attraverso dimensioni di modello da 124M a 2,7B parametri dimostrano miglioramenti consistenti.

Risultati Sperimentali

• Miglioramenti della Perplexity: Muown ottiene costantemente una perplexity inferiore rispetto a Muon, SOAP, AdamW e Lion su tutte le scale di modello testate (da 124M a 2,7B) e sui diversi tassi di apprendimento.
• Robustezza agli Iperparametri:
  • Tasso di Apprendimento: Muown amplia il plateau dei tassi di apprendimento quasi ottimali attraverso diverse larghezze di modello, riducendo la sensibilità alla sintonizzazione del tasso di apprendimento per scala.
  • Decadimento dei Pesi: Muown allevia la necessità di sintonizzare il decadimento dei pesi; funziona bene senza decadimento dei pesi ($\lambda=0$), mentre Muon richiede una sintonizzazione accurata (ottimo $\lambda=0.1$) per prevenire la divergenza o prestazioni subottimali.
• Stabilità: Muown previene la deriva sistematica della norma spettrale osservata in Muon. In regimi ad alto tasso di apprendimento (es. $4 \times 10^{-3}$), Muon spesso diverge, mentre Muown rimane stabile e vicino all'ottimo.
• Proprietà del Gradiente: L'analisi mostra che Muown mantiene un rango effettivo normalizzato più alto dei gradienti dei minibatch ed esibisce un coefficiente di rumore stocastico più basso durante tutto l'addestramento rispetto a Muon.
• Sovraccarico Computazionale: Il metodo introduce un sovraccarico trascurabile nel tempo per passo. Quando l'ottimizzatore è frammentato (sharded) tra i rank (ambiente distribuito), i tempi di passo per Muown e Muon diventano indistinguibili. Il sovraccarico di memoria è minimo (lineare nel numero di righe), aggiungendo solo pochi megabyte per livello.

Significato e Affermazioni

Il lavoro afferma che la deriva della norma spettrale osservata in Muon non è una proprietà irriducibile dell'ottimizzatore, ma una conseguenza della crescita incontrollata delle magnitudini delle righe. Isolando questo driver e promuovendolo a variabile esplicita dell'ottimizzatore, Muown ottiene:

• Convergenza Migliorata: Migliore perplexity e stabilità senza la necessità di un decadimento dei pesi aggressivo.
• Solidità Teorica: Tassi di convergenza ottimali in un framework consapevole della geometria che tiene esplicitamente conto della natura mista degli aggiornamenti.
• Efficienza Pratica: Una sostituzione diretta che corrisponde al tempo di esecuzione di Muon negli ambienti distribuiti offrendo al contempo una robustezza superiore alle scelte degli iperparametri.

Gli autori posizionano Muown come una soluzione che colma il divario tra i requisiti di stabilità teorica dell'apprendimento delle caratteristiche e le esigenze pratiche dell'addestramento di LLM su larga scala, offrendo un'alternativa più mirata al decadimento dei pesi disaccoppiato.