Technische Samenvatting: Muon leert robuustere en overdraagbare kenmerken dan Adam

Probleemstelling

Hoewel de Muon-optimizer (Momentum Orthogonalized by Newton-Schulz) recentelijk state-of-the-art efficiëntie heeft aangetoond bij het pretrainen van Large Language Models (LLMs) en visuele classificatieclassificaties—waarbij het vaak bijna twee keer zo snel traint als Adam en SGD—blijft de specifieke aard van de voordelen van Muon bij het leren van kenmerken (features) onduidelijk. Bestaande literatuur heeft het succes van Muon grotendeels verklaard vanuit optimalisatietheoretische lenzen, zoals steilste gradiëntafdaling onder spectrale normen, associatieve geheugendynamiek of het omgaan met heavy-tailed data. Dit onderzoek richt zich op de vraag: Leert Muon betere kenmerken dan Adam en SGD, en welke eigenschappen karakteriseren dit voordeel?

Methodologie

De auteurs hanteren een tweeledige aanpak die uitgebreide empirische evaluatie over visuele en talige domeinen combineert met theoretische analyse op een gestileerd classificatieprobleem.

Empirische Evaluatie

De studie evalueert de kwaliteit van kenmerken via twee primaire metrieken: robuustheid en overdraagbaarheid.

• Robuustheid: Geprepareerde modellen (ResNet-18, ViT-S, GPT-2, GPT-2 Medium) worden getest op gecorrumpeerde inputs. Visuele modellen worden geëvalueerd op ImageNet-C (15 corruptietypen), terwijl taalmodellen worden geëvalueerd op een nieuw geconstrueerde benchmark, FineWeb10B-C, die karakter-niveau typefouten (invoegingen, verwijderingen, substituties, wisselingen) toepast.
• Overdraagbaarheid: Geprepareerde kenmerken worden aangepast aan downstream-taken. Voor visuele taken worden lineaire classificators getraind op bevroren backbones over vier datasets (EuroSAT, Food-101, Flowers-102, Stanford Cars). Voor taalmodellen worden volledige modellen gefinetuned op instructie-volgende datasets (Alpaca, Dolly, WizardLM).
• Analyse van verborgen toestanden (Hidden-State Analysis): Om het mechanisme achter deze prestatiewinsten te begrijpen, analyseren de auteurs de laag-specifieke verborgen toestanden met behulp van:
  • Logit Margins: Getrainde lineaire probes (tuned-lens framework) decoderen de verborgen toestanden om de marge te berekenen tussen de correcte klasse/token en de sterkste afleider (distractor).
  • Spectrale Diversiteit: De effectieve rang en de Top-k energiefractie van de matrices van de verborgen toestanden worden berekend om de diversiteit en distributie van informatie over orthogonale richtingen te meten.
• Experimenteel Protocol: Alle vergelijkingen volgen een protocol met een gelijkwaardig budget waarbij architectuur, data en trainingsduur vaststaan, terwijl hyperparameters (leersnelheid, scheduler) onafhankelijk voor elke optimizer worden afgestemd om verschillen in representatie te isoleren van engineering-inspanning.

Theoretische Analyse

De auteurs bieden theoretische ondersteuning met behulp van een een-laags lineair classificatiemodel met multi-component kenmerken. Het datamodel gaat ervan uit dat kenmerken bestaan uit een gedeelde component op blokniveau (coarse class) en een klasse-specifieke component (fine-grained class). De analyse vergelijkt de trainingsdynamiek van Gradient Descent (GD), Adam en Muon onder infinitesimale stapgroottes en nul momentum, met de focus op de toestand waarin elke optimizer voor het eerst een specifieke verliesdrempel bereikt.

Belangrijkste Bijdragen en Resultaten

1. Empirisch Bewijs van Superieure Kenmerkkwaliteit

Het paper stelt vast dat Muon consistent kenmerken leert die robuuster en overdraagbaarder zijn dan die geleerd door Adam en SGD:

• Robuustheid: Muon-geprepareerde modellen behalen de hoogste gemiddelde nauwkeurigheid op ImageNet-C en de laagste perplexiteit op FineWeb10B-C over alle corruptietypen en architecturen heen. Opvallend genoeg presteert Muon beter dan Adam in vision transformers en taalmodellen, en beter dan SGD in taalmodellen.
• Overdraagbaarheid: Muon-geprepareerde kenmerken worden effectiever overgedragen naar downstream-taken. In visuele taken evenaart of overtreft Muon Adam en SGD op alle downstream-benchmarks. In taalmodellen bereiken Muon-geprepareerde modellen een lagere perplexiteit na instructie-tuning vergeleken met Adam-geprepareerde modellen.

2. Mechanismen van Verborgen Toestanden

De empirische voordelen zijn direct gekoppeld aan specifieke eigenschappen van de geleerde verborgen representaties:

• Grotere Logit Margins: Laag-specifieke probing onthult dat Muon aanzienlijk grotere logit margins oplevert over de meeste lagen vergeleken met Adam en SGD. Dit duidt op een bredere scheiding tussen correcte voorspellingen en afleiders, wat direct correleert met de geobserveerde robuustheid.
• Hogere Effectieve Rang: Spectrale analyse laat zien dat Muon een hogere effectieve rang en een lagere Top-10 energiefractie bereikt over de lagen heen. Dit suggereert dat de representaties van Muon informatie gelijkmatiger verspreiden over orthogonale richtingen, wat een rijkere feature-subspace creëert die betere adaptatie aan downstream-taken faciliteert.

3. Theoretische Garanties

In de gestileerde multi-component kenmerksetting bewijzen de auteurs dat de spectrale normalisatie van Muon de onbalans in singuliere waarden in de gradiëntmatrix wegneemt.

• Margin Voordeel: Muon bereikt een strikt grotere classificatiemarge dan Adam en GD bij een gematchte trainingsverlies (buiten een exponentieel kleine terminal-loss window). Dit komt omdat de spectrale normalisatie van Muon de bijdrage van orthogonale feature-richtingen balanceert, terwijl de coördinaat-gewijze normalisatie van Adam en de ruwe gradiëntstroom van GD leiden tot een overmatige afhankelijkheid van specifieke componenten.
• Effectieve Rang Voordeel: Muon bereikt bewezen een strikt hogere effectieve rang van kenmerken dan Adam en GD bij een gematcht verlies, wat de empirische bevinding van meer diverse representaties ondersteunt.

Betekenis en Claims

Het paper claimt de voordelen van Muon te "demystificeren" door de focus te verschuiven van optimalisatie-efficiëntie naar kenmerkkwaliteit. De auteurs argumenteren dat het succes van Muon niet louter een resultaat is van snellere convergentie, maar voortkomt uit het vermogen om representaties te leren die inherent robuuster zijn tegen input-perturbaties en beter aanpasbaar zijn aan nieuwe taken.

De betekenis van dit werk ligt in:

1. Het Verbinden van Optimalisatie met Representatie: Het legt een concreet verband tussen de matrix-bewuste update-regel van Muon en de geometrische eigenschappen van geleerde kenmerken (logit margins en effectieve rang).
2. Theoretische Validatie: Het biedt het eerste theoretische bewijs dat Muon systematisch verschillende en superieure kenmerkkwaliteiten induceert (grotere marges, hogere effectieve rang) vergeleken met Adam en GD onder condities van gematcht verlies.
3. Praktische Implicaties: De bevindingen suggereren dat voor toepassingen die robuustheid vereisen tegen distributieverschuivingen of sterke overdraagbaarheid naar downstream-taken, Muon een superieure keuze is boven standaard adaptieve optimizers zoals Adam, zelfs wanneer de trainingsbudgetten gelijk zijn.

De auteurs merken bescheiden op dat hun huidige analyse zich richt op LLM's en visuele classificatieclassificaties, waarbij de prestaties van Muon op andere modelfamilies (bijv. diffusiemodellen) een richting voor toekomstig werk laten.