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

Technische Samenvatting: Muon leert gebalanceerde oplossingen in matrixfactorisatie zonder trage saddle-to-saddle dynamiek

Probleemstelling
Het paper onderzoekt de leerdynamiek van de Muon-optimizer in de context van matrixfactorisatie, gedefinieerd als het minimaliseren van $L(P, Q) = \frac{1}{2}\|M^* - PQ^\top\|_F^2$. Matrixfactorisatie dient als een minimaal model voor het bestuderen van niet-lineaire parameterdynamiek en representatieleer, waarbij gedragingen worden vertoond die relevant zijn voor deep learning, zoals low-rank bias en representatie-uitlijning. De auteurs vergelijken specifiek Muon met standaard Gradient Descent (GD) om te begrijpen hoe de unieke update-regel van Muon het traject van optimalisatie verandert, met name met betrekking tot de snelheid van convergentie, stabiliteit en de aard van de gevonden oplossingen vanuit een kleine willekeurige initialisatie.

Methodologie
De studie richt zich op de Muon update-regel, die de gradiënt orthogonaliseert voordat de learning rate wordt toegepast. Voor een gewichtsmatrix $W$ is de update $\Delta W = -\eta \cdot \widehat{\nabla_W L}$, waarbij $\widehat{G} = UV^\top$ de poolontbinding (orthogonalisering) is van de gradiënt $G=USV^\top$. Dit zorgt ervoor dat alle updates (semi-)orthogonale matrices zijn, geschaald met $\eta$.

De auteurs analyseren de dynamiek via drie lenzen:

1. Scalaire Reductie: Ervan uitgaande dat gewichten uitgelijnd zijn, ontkoppelt het probleem zich in onafhankelijke scalaire problemen waarbij de residu $r = m^* - pq$ de update dicteert.
2. Uitlijningsdynamiek: De auteurs leiden theoretische snelheden af voor hoe de singuliere vectoren van $P$ en $Q$ met elkaar uitlijnen (interne uitlijning) en met de doelmatrix $M^*$ (doel-uitlijning) onder verschillende parameterregimes (onder-geparameteriseerd, exact-geparameteriseerd en over-geparameteriseerd).
3. Geconserveerde Grootheden: Het paper identificeert specifieke matrix-grootheden die geconserveerd blijven tijdens de flow van Muon versus GD.

Belangrijkste Bijdragen en Resultaten

1. Vermijden van trage saddle-to-saddle dynamiek
Standaard Gradient Descent vanuit een kleine initialisatie vertoont "saddle-to-saddle" dynamiek. Singuliere waarden groeien multiplicatief, wat betekent dat kleinere modi langzamer groeien dan grotere. De optimizer blijft hangen nabij low-rank saddles, waarbij modi sequentieel worden geleerd (stapsgewijze leerproces).
In tegenstelling hiertoe vertoont Muon additieve dynamiek. Zodra uitlijning plaatsvindt, groeien alle singuliere waarden van de gewichtsmatrices met een uniforme snelheid ($\eta$ per stap), ongeacht hun huidige grootte. Gevolgelijk convergeren kleinere modi voordat grotere modi, en leert de optimizer alle top-modi van $M^*$ simultaan zonder tussenliggende zadelpunten.

2. Stabiliteit voorbij kritieke learning rates
Gradient Descent is beperkt door een kritieke learning rate drempel bepaald door de lokale loss-scherpte (specifiek $2/\lambda_{\max}(\nabla^2 L)$). Het overschrijden hiervan veroorzaakt divergentie, en de convergentiesnelheid wordt beperkt door de conditiegetal van het probleem.
Muon blijft stabiel zelfs wanneer de learning rate deze kritieke drempel overschrijdt. Omdat de update-grootte wordt vastgesteld door de orthogonaliseringsstap, kunnen de iteraties niet divergeren. In plaats daarvan, bij grote learning rates, procedeert de optimalisatie snel naar de oplossing-manifold en oscilleert daar. Dit ontkoppelt de learning rate van het conditiegetal, wat snelle convergentie mogelijk maakt via exponentiële (geometrische) learning rate annealing.

3. Distinctieve Geconserveerde Grootheden en Gebalanceerde Oplossingen
Beide optimizers vinden "gebalanceerde oplossingen" (waarbij $P^\top P = Q^\top Q$) vanuit een verdwijnende initialisatie, maar ze conserveren verschillende grootheden:

• Gradient Flow: Conserveert de matrix-grootheid $\Delta_2 = P^\top P - Q^\top Q$. De flow-lijnen in de parameterruimte zijn hyperbolisch.
• Muon Flow: Conserveert de matrix-grootheid $\Delta_1 = \sqrt{P^\top P} - \sqrt{Q^\top Q}$. De flow-lijnen zijn affien (parallelle stralen).
  De auteurs leiden af dat de Muon flow wordt gegenereerd door een Hamiltoniaan $H(P, Q) = \text{Tr}(\Delta_1) = \|P\|_* - \|Q\|_*$ (het verschil in nucleaire normen). Deze conserveringswet garandeert dat als gewichten gebalanceerd starten, ze gebalanceerd blijven, wat leidt tot de gebalanceerde oplossing.

4. Uitlijningsdynamiek en Snelheden
Het paper leidt uitlijningssnelheden af voor de "silent alignment" fase (waarbij singuliere vectoren uitlijnen terwijl de loss vlak is):

• Onder-geparameteriseerd ($d < n$): Interne uitlijning is triviaal; doel-uitlijning vervalt polynomiaal als $(t/\alpha_{\text{eff}})^{-2(1-s^{-1})}$, waarbij $s$ de relatieve spectrale gap is.
• Over-geparameteriseerd ($d > n$): Doel-uitlijning is triviaal; interne misuitlijning vervalt snel als $(t/\alpha_{\text{eff}})^{-4}$.
• Algemeen: Uitlijning vindt eerder plaats met kleinere initialisatieschalen. De auteurs tonen aan dat deze theoretische snelheden empirische uitlijning in algemene settings voorspellen.

5. Spiked Learning Rate Schedule
Door gebruik te maken van het begrip van uitlijningsdynamiek, construeren de auteurs een specifieke learning rate schedule die bijna perfecte uitlijning bereikt in exact twee optimalisatiestappen (voor $d \ge n$):

1. Stap 1: Stel $\eta^{(0)} = \alpha$ (initialisatieschaal) in om $P$ met $Q$ uit te lijnen.
2. Stap 2: Stel een grote spike $\eta^{(1)} \gg \eta^{(0)}$ in om $P$ simultaan met de doelmatrix $M^*$ uit te lijnen en de gewichten naar de oplossing-manifold te drijven.
3. Annealing: Voer geometrische annealing uit op de learning rate om lineaire convergentie te bereiken.
   Dit schema demonstreert dat theoretische inzichten in de flow vertaald kunnen worden naar praktische optimalisatiestrategieën, hoewel de auteurs opmerken dat dit een proof-of-concept is voor specifieke regimes.

Betekenis en Claims
Het paper claimt dat Muon de optimalisatielandschap van matrixfactorisatie fundamenteel verandert door multiplicatieve, saddle-rijke dynamiek te vervangen door uniforme, additieve groei. Dit stelt Muon in staat om:

• Alle modi van de doelmatrix met dezelfde snelheid te leren, waardoor het trage, sequentiële leerproces van GD wordt vermeden.
• Stabiel te opereren met learning rates die ver boven de kritieke drempels liggen die GD beperken, waardoor de learning rate wordt bevrijd van het conditiegetal van het probleem.
• Gebalanceerde oplossingen te bereiken via een distinctieve conserveringswet ($\sqrt{P^\top P} - \sqrt{Q^\top Q}$) die gewichtsbalans handhaaft zonder de hyperbolische geometrie van GD te vereisen.

De auteurs benadrukken dat hoewel hun analyse beperkt is tot ongewijzigde spectrale gradient descent (waarbij momentum, weight decay en Newton-Schulz iteraties worden genegeerd) en eenvoudige matrixfactorisatie, het mechanistische begrip van uitlijning en uniforme groei een fundament biedt voor het begrijpen van het succes van Muon in grotere, complexere modellen. Het werk dient als een theoretische verklaring voor waarom Muon training versnelt in slecht geconditioneerde of heterogene spectrale settings.