Technische Samenvatting: Herbeoordeling van Muon voor Matrixfactorisatie

Probleemstelling

Muon (MomentUm Orthogonalized by Newton–Schulz) is onlangs opgekomen als een krachtige optimizer voor grootschalige deep learning, met name bij het trainen van large language models (LLMs), waar het naar verluidt beter presteert dan standaard adaptieve methoden zoals Adam en AdamW. De theoretische literatuur interpreteert Muon als het uitvoeren van de steilste daling onder de spectrale norm, waarbij het gebruikmaakt van benaderde orthogonalisatie om gradiënt-updates te hervormen.

Echter, er blijft een kritieke kloof bestaan: het is onduidelijk of het empirische succes van Muon voortkomt uit zijn intrinsieke update-regel of dat het een artefact is van de specifieke schaal, architectuur en datadistributies van moderne diepe netwerken. Bestaande theoretische analyses leunen vaak op geïdealiseerde settings of lokale eigenschappen, en slagen er niet in om concrete condities te specificeren waaronder Muon theoretisch of empirisch zou moeten presteren beter dan klassieke optimizers.

Dit artikel adresseert deze onzekerheid door de optimizer te isoleren van verwarrende factoren (architectuur, schaal, data) door Muon te evalueren op low-rank matrixfactorisatie. Dit probleem is gekozen vanwege de eenvoud, het goed begrepen optimalisatielandschap en de expliciete spectrale structuur, wat een systematische en uitputtende evaluatie mogelijk maakt van hoe de spectrale aard van Muon interageert met kromming en singular-waarde structuren.

Methodologie

De auteurs voeren een gecontroleerde, systematische vergelijking uit tussen Muon en adaptieve baselines (AdamW, plain Gradient Descent, en SignGD) over vijf verschillende matrix-herstel scenario's. Het evaluatieprotocol is ontworpen om bias door hyperparameterselectie te elimineren, wat de auteurs identificeren als een primaire bron van misleidende conclusies in eerder werk.

Experimenteel Ontwerp

1. Probleemvarianten:

   • Symmetrische Low-Rank Factorisatie: Het herstellen van een rank-$r$ positief semidefiniete matrix $M^\star$ uit $UU^\top$.
   • Matrix Completion: Het herstellen van een matrix uit een subset van geobserveerde entries (20% observatieratio).
   • Non-negative Matrix Factorization (NMF): Een geconstrainde factorisatie waarbij zowel de doelmatrix als de factoren niet-negatief zijn.
   • Tensor-Train Factorisatie: Een bilaterale factorisatie-setup toegepast op symmetrische PSD-doelen, dienend als een minimaal model voor product-parameterisaties in deep learning.
   • Gaussian-Kernel Factorisatie: Een glad, goed geconditioneerd landschap met een RBF-kernmatrix.

2. Gecontroleerde Variabelen:

   • Conditionering ($\kappa$): Experimenten scannen condition numbers van goed geconditioneerd ($\kappa=1$) tot ernstig slecht geconditioneerd ($\kappa=625$ of $10^4$).
   • Spectrale Vorm: De distributie van eigenwaarden wordt gevarieerd (uniform, vervallend, heavy-tailed, geclusterd) terwijl de condition number constant wordt gehouden om de gevoeligheid voor spectrale dichtheid te testen.
   • Parameterisatie: Experimenten dekken onder-geparameteriseerde, exact geparameteriseerde en over-geparameteriseerde regimes.

3. Evaluatieprotocol:

   • Hyperparameter Tuning: In plaats van een enkele "standaard" learning rate, voeren de auteurs uitputtende learning rate sweeps uit (logaritmische grids) voor elke optimizer en elke probleemconfiguratie.
   • Metriek: Herstel wordt gemeten via de Normalized Mean-Squared Error (NMSE).
   • Vergelijking: Optimizers worden vergeleken op basis van hun "stabiele learning-rate bereik" (de band van stapgroottes waarbij convergentie naar de loss floor optreedt) en hun best-getunede finale loss.

Belangrijkste Bijdragen en Bevindingen

1. Hyperparameter Gevoeligheid en Tuning

Het artikel demonstreert dat de rangorde van optimizers zeer gevoelig is voor de selectie van de learning rate.

• Default vs. Getuned: Rangordes afgeleid van een enkele standaard learning rate (bijv. $10^{-3}$) spreken de rangordes die verkregen zijn na individuele tuning vaak tegen.
• Geen Universele Dominantie: Wanneer elke methode individueel wordt getuned, domineert geen enkele optimizer alle settings.
  • AdamW en GD evenaren of overtreffen vaak Muon bij standaard low-rank factorisatie en matrix completion taken.
  • Muon vertoont een consistent voordeel, primair in Non-negative Matrix Factorization (NMF) en specifieke, sterk slecht geconditioneerde completion problemen.

2. De Rol van Conditionering en Spectrale Structuur

• Conditionering Herordent Methoden: Naarmate de condition number ($\kappa$) toeneemt, verschuift de optimale learning rate en verandert de relatieve prestatie van optimizers. Bijvoorbeeld: GD presteert goed in goed geconditioneerde regimes, maar faalt in extreme ill-conditioning, waar AdamW uitblinkt.
• Spectrale Vorm Is Relevant: Zelfs bij een vaste, extreme condition number ($\kappa=10^4$), heeft de interne distributie van eigenwaarden een significante impact op de prestaties. AdamW past zich effectief aan heavy-tailed en geclusterde spectra aan, terwijl Muon onderpresteert bij de meeste spectrale vormen in deze extreme condities.

3. Specifieke Prestaties in NMF

Muon behoudt een duidelijk voordeel in NMF-taken. De auteurs suggereren dat dit komt doordat de georthogonaliseerde updates van Muon redundante factoren ontmoedigen en meer diverse representaties bevorderen, een voordeel dat niet wordt waargenomen in ongestelde of symmetrische factorisatietaken.

4. Goed Gedragde Landschappen

Bij de Gaussian-kernel factorisatie taak (een glad, goed geconditioneerd landschap), biedt het complexe orthogonalisatie-mechanisme van Muon geen voordeel ten opzichte van een goed getunede AdamW. Beiden convergeren naar dezelfde globale floor, wat bevestigt dat de voordelen van Muon niet universeel zijn maar contextafhankelijk.

Betekenis en Claims

Het artikel betoogt dat de vermeende superioriteit van Muon in grootschalige deep learning mogelijk verward wordt met andere factoren zoals model schaal, architectuur en datakenmerken. Door de optimizer te isoleren op gecontroleerde problemen, biedt het werk een genuanceerder beeld van de capaciteiten van Muon:

• Contextafhankelijke Nuttigheid: Muon is geen universeel superieure optimizer. De voordelen zijn verbonden aan specifieke probleemstructuren, met name die met non-negativiteitsrestricties of specifieke spectrale degeneraties die in NMF worden gevonden.
• Evaluatiestandaarden: Het werk onderstreept de noodzaak om moderne optimizers te evalueren over getunede hyperparameter ranges en diverse conditioneringsregimes, in plaats van te vertrouwen op standaardinstellingen of geïsoleerde end-to-end benchmarks.
• Theoretische Implicaties: De resultaten dagen de aanname uit dat de spectrale-norm steile daling (de theoretische interpretatie van Muon) inherent superieur is aan coördinaat-wijze adaptiviteit (AdamW) in alle gestructureerde settings. In plaats daarvan hangt de "beste" aanpak sterk af van de specifieke geometrische en spectrale kenmerken van de objectieve functie.

Concluderend stellen de auteurs niet dat Muon ineffectief is, maar dat de voordelen ervan probleemafhankelijk zijn. De studie pleit voor een striktere, gecontroleerde benchmarking-aanpak om te begrijpen wanneer en waarom spectrum-bewuste orthogonalisatie de optimalisatie helpt.