Muon learns balanced solutions in matrix factorization without slow saddle-to-saddle dynamics

Sintesi Tecnica: Muon apprende soluzioni bilanciate nella fattorizzazione di matrici senza la lenta dinamica da sella a sella

Problema
Il documento investiga le dinamiche di apprendimento dell'ottimizzatore Muon nel contesto della fattorizzazione di matrici, definita come la minimizzazione di $L(P, Q) = \frac{1}{2}\|M^* - PQ^\top\|_F^2$. La fattorizzazione di matrici funge da modello minimo per studiare le dinamiche dei parametri non lineari e l'apprendimento delle rappresentazioni, esibendo comportamenti rilevanti per il deep learning come il bias a basso rango (low-rank bias) e l'allineamento delle rappresentazioni. Gli autori confrontano specificamente Muon con il classico Gradient Descent (GD) per comprendere come la regola di aggiornamento unica di Muon alteri la traiettoria dell'ottimizzazione, in particolare per quanto riguarda la velocità di convergenza, la stabilità e la natura delle soluzioni trovate partendo da una piccola inizializzazione casuale.

Metodologia
Lo studio si concentra sulla regola di aggiornamento di Muon, che ortogonalizza il gradiente prima di applicare il tasso di apprendimento (learning rate). Per una matrice di pesi $W$, l'aggiornamento è $\Delta W = -\eta \cdot \widehat{\nabla_W L}$, dove $\widehat{G} = UV^\top$ è la decomposizione polare (ortogonalizzazione) del gradiente $G=USV^\top$. Ciò assicura che tutti gli aggiornamenti siano matrici (semi-)ortogonali scalate per $\eta$.

Gli autori analizzano le dinamiche attraverso tre lenti:

1. Riduzione Scalare: Assumendo che i pesi siano allineati, il problema si decompone in problemi scalari indipendenti dove il residuo $r = m^* - pq$ determina l'aggiornamento.
2. Dinamica di Allineamento: Gli autori derivano i tassi teorici di come i vettori singolari di $P$ e $Q$ si allineano tra loro (allineamento interno) e con la matrice target $M^*$ (allineamento con il target) sotto diversi regimi di parametri (sotto-parametrizzato, esattamente parametrizzato e sovra-parametrizzato).
3. Quantità Conservate: Il documento identifica specifiche quantità matriciali conservate durante il flusso di Muon rispetto a GD.

Contributi Chiave e Risultati

1. Evitare la lenta dinamica da sella a sella
Il classico Gradient Descent, partendo da una piccola inizializzazione, esibisce una dinamica "da sella a sella". I valori singolari crescono moltiplicativamente, il che significa che i modi più piccoli crescono più lentamente di quelli più grandi. L'ottimizzatore ristagna vicino a selle a basso rango, apprendendo i modi in modo sequenziale (apprendimento a step).
Al contrario, Muon esibisce una dinamica additiva. Una volta avvenuto l'allineamento, tutti i valori singolari delle matrici di peso crescono a un tasso uniforme ($\eta$ per step), indipendentemente dalla loro grandezza attuale. Di conseguenza, i modi più piccoli convergono prima di quelli più grandi, e l'ottimizzatore apprende tutti i modi superiori di $M^*$ simultaneamente, senza punti di sella intermedi.

2. Stabilità oltre i tassi di apprendimento critici
Il Gradient Descent è vincolato da una soglia critica del tasso di apprendimento determinata dalla nitidezza locale della perdita (specificamente $2/\lambda_{\max}(\nabla^2 L)$). Superare tale soglia causa divergenza, e il tasso di convergenza è limitato dal numero di condizionamento del problema.
Muon rimane stabile anche quando il tasso di apprendimento supera questa soglia critica. Poiché l'entità dell'aggiornamento è fissata dal passaggio di ortogonalizzazione, le iterazioni non possono divergere. Invece, con tassi di apprendimento elevati, l'ottimizzazione procede rapidamente verso il manifold della soluzione e oscilla lì. Ciò scinde il tasso di apprendimento dal numero di condizionamento, consentendo una convergenza rapida tramite un annealing (attenuazione) geometrico del tasso di apprendimento.

3. Quantità conservate distinte e soluzioni bilanciate
Entrambi gli ottimizzatori trovano "soluzioni bilanciate" (dove $P^\top P = Q^\top Q$) da un'inizializzazione nulla, ma conservano quantità diverse:

• Flusso del Gradiente: Conserva la quantità matriciale $\Delta_2 = P^\top P - Q^\top Q$. Le linee di flusso nello spazio dei parametri sono iperboliche.
• Flusso di Muon: Conserva la quantità matriciale $\Delta_1 = \sqrt{P^\top P} - \sqrt{Q^\top Q}$. Le linee di flusso sono affini (raggi paralleli).
  Gli autori derivano che il flusso di Muon è generato da un Hamiltoniano $H(P, Q) = \text{Tr}(\Delta_1) = \|P\|_* - \|Q\|_*$ (la differenza nelle norme nucleari). Questa legge di conservazione assicura che se i pesi partono bilanciati, rimangano bilanciati, portando alla soluzione bilanciata.

4. Dinamiche di allineamento e tassi
Il documento deriva i tassi di allineamento per la fase di "allineamento silenzioso" (dove i vettori singolari si allineano mentre la perdita è piatta):

• Sotto-parametrizzato ($d < n$): L'allineamento interno è banale; l'allineamento con il target decade polinomialmente come $(t/\alpha_{\text{eff}})^{-2(1-s^{-1})}$, dove $s$ è il gap spettrale relativo.
• Sovra-parametrizzato ($d > n$): L'allineamento con il target è banale; il disallineamento interno decade rapidamente come $(t/\alpha_{\text{eff}})^{-4}$.
• Generale: L'allineamento avviene prima con scale di inizializzazione più piccole. Gli autori mostrano che questi tassi teorici predicono l'allineamento empirico in contesti generali.

5. Programma di apprendimento a picco (Spiked Learning Rate Schedule)
Sfruttando la comprensione delle dinamiche di allineamento, gli autori costruiscono un programma specifico del tasso di apprendimento che raggiunge un allineamento quasi perfetto in esattamente due step di ottimizzazione (per $d \ge n$):

1. Step 1: Impostare $\eta^{(0)} = \alpha$ (scala di inizializzazione) per allineare $P$ con $Q$.
2. Step 2: Impostare un picco elevato $\eta^{(1)} \gg \eta^{(0)}$ per allineare simultaneamente il modello con il target $M^*$ e spingere i pesi verso il manifold della soluzione.
3. Annealing: Effettuare un annealing geometrico del tasso di apprendimento per ottenere una convergenza lineare.
   Questo programma dimostra che le intuizioni teoriche sul flusso possono essere tradotte in strategie di ottimizzazione pratiche, sebbene gli autori la considerino una prova di concetto per regimi specifici.

Significato e Rivendicazioni
Il documento sostiene che Muon altera fondamentalmente il panorama di ottimizzazione della fattorizzazione di matrici, sostituendo le dinamiche moltiplicative cariche di selle con una crescita additiva uniforme. Ciò consente a Muon di:

• Apprendere tutti i modi della matrice target allo stesso tasso, evitando l'apprendimento sequenziale e lento tipico del GD.
• Operare stabilmente con tassi di apprendimento che superano di gran lunga le soglie critiche che vincolano il GD, liberando il tasso di apprendimento dal numero di condizionamento del problema.
• Raggiungere soluzioni bilanciate tramite una distinta legge di conservazione ($\sqrt{P^\top P} - \sqrt{Q^\top Q}$) che mantiene l'equilibrio dei pesi senza richiedere la geometria iperbolica del GD.

Gli autori sottolineano che, sebbene la loro analisi sia limitata a un gradiente spettrale non modificato (ignorando momentum, weight decay e iterazioni di Newton-Schulz) e alla semplice fattorizzazione di matrici, la comprensione meccanicistica di allineamento e crescita uniforme fornisce una base per comprendere il successo di Muon in modelli più grandi e complessi. Il lavoro serve come spiegazione teorica del perché Muon acceleri l'addestramento in contesti spettrali mal condizionati o eterogenei.