To Use or not to Use Muon: How Simplicity Bias in Optimizers Matters

Deze paper waarschuwt dat de optimizer Muon, ondanks zijn superieure trainsnelheid, de natuurlijke 'simplicity bias' van methoden zoals SGD verwijdert, wat kan leiden tot het leren van spurious features en een verminderd vermogen om onderliggende structuren te ontdekken.

De kern

Titel: Snelheid versus Slimheid: Waarom de nieuwe optimisator 'Muon' soms te snel gaat

Stel je voor dat je een heel groot labyrint moet doorkruisen om de schat te vinden. In de wereld van kunstmatige intelligentie (AI) is dit labyrint de "foutenkaart" van een computermodel. De schat is de perfecte oplossing.

Jarenlang was SGD (Stochastic Gradient Descent) de standaard. Het is als een wandelaar die heel voorzichtig en systematisch door het labyrint loopt. Hij bekijkt elke hoek, loopt soms vast in een klein dal (een "zadel"), en moet daar even uitklimmen voordat hij verder kan. Dit proces is traag, maar het zorgt ervoor dat hij de eenvoudigste route vindt. Hij leert stap voor stap, eerst de grote paden, dan de kleinere.

Onlangs is er een nieuwe, razendsnelle runner genaamd Muon op de markt gekomen. Muon is als een helikopter die boven het labyrint vliegt. Hij ziet direct waar de schat ligt en vliegt er rechtstreeks naartoe. Hij slaat alle kleine dalen en obstakels over. Hierdoor is hij veel sneller dan de wandelaar.

Maar wat is de prijs van die snelheid?

Deze paper (van Sara Dragutinović en Rajesh Ranganath) waarschuwt: "Snelheid is geweldig, maar Muon mist een belangrijk voordeel van de wandelaar: de eenvoudigheid."

Hier is de uitleg in drie simpele analogieën:

1. Het "Stap-voor-stap" versus "Alles-terzelfdertijd" probleem

• De Wandelaar (SGD): Stel je voor dat je een taal moet leren. De wandelaar leert eerst de basiswoorden, dan de zinnen, en pas daarna de complexe grammatica. Hij bouwt zijn kennis op. Dit heet een "simpliciteitsbias" (voorkeur voor eenvoud). Omdat hij eerst de basis goed heeft, begrijpt hij de logica van de taal beter.
• De Helikopter (Muon): Muon leert alles tegelijk. Hij leert de basiswoorden, de zinnen én de grammatica in één klap. Dit gaat razendsnel. Maar omdat hij alles door elkaar leert, mist hij de kans om te zien dat de basisregels eigenlijk hetzelfde zijn voor verschillende situaties. Hij leert de regels niet, hij leert alleen de oplossingen voor de specifieke vragen die hij ziet.

2. De "Gemeenschappelijke Structuur" (Het Routeer-Experiment)

De auteurs testten dit met een spelletje. Stel je voor dat je 7 verschillende talen spreekt, maar ze hebben allemaal dezelfde basiswoorden (bijv. "1, 2, 3, 4").

• SGD: Omdat hij stap voor stap leert, ziet hij dat de basiswoorden in alle talen hetzelfde zijn. Hij leert één "gemeenschappelijke hersenstructuur" die voor alle talen werkt. Als hij later een nieuwe taal ziet, kan hij die direct begrijpen.
• Muon: Omdat hij alles tegelijk leert, ziet hij alleen de specifieke woorden voor elke taal apart. Hij onthoudt: "In taal A betekent dit woord X, in taal B betekent dit woord Y." Hij heeft geen idee dat ze eigenlijk hetzelfde zijn. Als hij een nieuwe taal ziet, faalt hij, omdat hij geen algemene regel heeft geleerd, maar alleen uit het hoofd heeft geleerd wat hij al gezien had.

Conclusie: Muon is goed voor het onthouden van feiten, maar slecht voor het begrijpen van de onderliggende logica.

3. De "Valse Vriend" (Spurious Features)

Stel je voor dat je een computer wilt leren om foto's van katten en honden te herkennen.

• De situatie: In de trainingsfoto's staan alle katten op een groene achtergrond en alle honden op een blauwe achtergrond. De echte kenmerken (oren, staart) zijn er ook, maar de achtergrond is heel opvallend.
• SGD (De voorzichtige): Omdat hij eerst de belangrijkste dingen leert, kijkt hij eerst naar de dieren zelf (oren, staart). Pas heel laat, als hij al bijna klaar is, ziet hij de achtergrond. Als je hem dan een foto geeft met een kat op een blauwe achtergrond, herkent hij de kat nog steeds, omdat hij de kat zelf heeft geleerd.
• Muon (De snelle): Omdat hij alles tegelijk leert, pakt hij direct de makkelijkste hint: de achtergrond. Hij leert: "Groen = Kat, Blauw = Hond". Hij is super snel in de test, maar faalt volledig als je hem een kat op een blauwe achtergrond laat zien. Hij heeft de "valse vriend" (de achtergrond) geleerd in plaats van de waarheid.

Wat betekent dit voor de toekomst?

De paper zegt niet dat Muon slecht is. Integendeel, Muon is vaak sneller en werkt goed in veel situaties. Maar de boodschap is: We moeten niet blindelings kiezen voor de snelste optie.

• Als je een taak hebt waarbij je snel veel feiten moet onthouden (zoals het onthouden van biografische gegevens), is Muon misschien superieur.
• Maar als je een taak hebt waarbij je de onderliggende logica moet begrijpen (zoals wiskunde, medische diagnoses, of het herkennen van objecten in nieuwe situaties), dan is de "trage" SGD soms beter omdat hij de eenvoudige, algemene regels vindt.

De kernboodschap:
Bij het ontwikkelen van nieuwe AI-tools moeten we niet alleen kijken naar "hoe snel is het?", maar ook naar "welke bias (voorkeur) introduceert deze tool?". Snelheid is geweldig, maar als we te snel gaan, kunnen we de diepere, slimme oplossingen missen en in plaats daarvan alleen maar oppervlakkige patronen onthouden.

Technische samenvatting

Probleemstelling

Hoewel de Muon-optimizer (MomentUm Orthogonalized by Newton-Schulz) recentelijk enorme populariteit heeft verworven vanwege zijn superieure trainsnelheid (bijvoorbeeld in de NanoGPT Speedrun en NanoChat), is er weinig inzicht in de inductieve bias die deze optimizer introduceert.

• Huidige situatie: Optimizers worden vaak geselecteerd op basis van convergentiesnelheid (wallclock time).
• Het probleem: Het is onbekend welke trajecten Muon door de verlieslandschap (loss landscape) aflegt en welke consequenties dit heeft voor de uiteindelijke oplossing.
• De kernvraag: Verliest Muon een cruciale "simplicity bias" (eenvoudsbias) die traditionele methoden zoals Stochastic Gradient Descent (SGD) van nature behouden, en leidt dit tot suboptimale generalisatie of het leren van spurious (schijn)correlaties?

Methodologie

De auteurs analyseren Muon theoretisch en empirisch door de volgende stappen te nemen:

1. Theoretische Vereenvoudiging (Spectral GD):

   • Om de wiskundige analyse haalbaar te maken, introduceren ze Spectral Gradient Descent (Spectral GD). Dit is een vereenvoudigde versie van Muon die exacte Singular Value Decomposition (SVD) gebruikt in plaats van de Newton-Schulz-iteraties, en geen momentum toepast.
   • Ze analyseren diepe lineaire netwerken (2 lagen), een standaardkader voor het bestuderen van leer-dynamica.
   • Ze vergelijken de dynamica van Gradient Flow (GD) met die van Spectral GD.

2. Theoretische Analyse:

   • Ze onderzoeken hoe de rang (rank) van de oplossing evolueert tijdens het trainen.
   • Ze analyseren de leertrajecten in termen van singuliere waarden van de productmatrix $W = VU$.

3. Empirische Validatie:

   • Experiment 1 (Gedeelde Representaties): Een "routing"-taak waarbij een model moet leren een onderliggende structuur te generaliseren over verschillende input-domeinen (multi-modaliteit), maar niet alle input-output paren tijdens het trainen ziet.
   • Experiment 2 (Spurious Features): Een MNIST-variant waarbij een spurious feature (een specifiek pixelintensiteit per klasse) aanwezig is. Het doel is te zien welk optimizer het eerst leert vertrouwen op de echte cijfers versus de spurious feature.

Kernbijdragen en Theoretische Bevindingen

1. Verlies van de "Simplicity Bias" bij Muon/Spectral GD

• SGD (Gradient Descent): Volgt een sequentiële leertraject. Het leert eerst de dominantste singuliere waarden (hoogste rang) en bouwt de complexiteit van de oplossing stap voor stap op. Dit wordt "saddle-to-saddle" dynamica genoemd, waarbij het model door verschillende zadelpunten beweegt. Dit gedrag fungeert als een impliciete regularisatie die de oplossing simpel houdt totdat het noodzakelijk is om complexer te worden.
• Muon (Spectral GD): Loopt gelijktijdig alle singuliere waarden bij. Het leert alle componenten (van dominant tot minder dominant) tegelijkertijd en met dezelfde snelheid.
• Gevolg: Muon mist de "simplicity bias". Het negeert de natuurlijke hiërarchie in de data en leert alles in één keer, wat leidt tot een snellere convergentie maar een gebrek aan structurele ordening.

2. Gevolgen voor Generalisatie

Het verlies van deze bias heeft twee belangrijke negatieve gevolgen:

• Moeilijkheid om gedeelde structuren te vinden: In plaats van een onderliggende, gedeelde representatie te leren die over verschillende domeinen werkt, neigt Muon naar memorisatie van specifieke input-output paren.
• Gevoeligheid voor spurious correlaties: Omdat Muon alle componenten tegelijk leert, kan het eerder gaan vertrouwen op schijnbare patronen (spurious features) als deze dominant zijn in de data, terwijl SGD eerst de "echte" dominante structuur leert.

Resultaten uit Experimenten

1. Routing-taak (Gedeelde Representaties):

   • SGD: Leert de onderliggende mapping (de "ware" structuur) en generaliseert perfect naar ongezette input-output paren. De oplossing convergeert naar een laag-rang structuur (rang = 4, gelijk aan de dimensie van de taak).
   • Spectral GD (Muon): Bereikt een perfecte loss op de trainingsdata, maar faalt bij generalisatie. Het leert de data te memoriseren in plaats van de structuur te begrijpen. De oplossing heeft een veel hogere effectieve rang en een zwaar-tail spectrum, wat duidt op overfitting op specifieke patronen in plaats van het leren van algemene regels.

2. Spurious Features (MNIST):

   • SGD: Blijft langer vertrouwen op de echte cijferstructuur. Alleen als de spurious feature extreem dominant wordt, schakelt SGD over. Dit resulteert in een hogere piek-accuraatheid op een validatiedataset zonder spurious features (bij vroege stop).
   • Muon: Loopt de echte structuur en de spurious feature tegelijkertijd bij. Het schakelt eerder over naar het vertrouwen op de spurious feature, wat leidt tot een snellere degradatie van de prestaties op schone data.
   • Conclusie: In scenario's waar de "echte" feature dominant is, presteert SGD beter omdat het de complexiteit stap voor stap opbouwt. Muon is sneller, maar minder robuust tegenover veranderingen in de data-distributie.

Significantie en Conclusie

• Snelheid vs. Kwaliteit: De paper waarschuwt dat de enorme snelheidswinst van Muon een prijs heeft: het verlies van de inductieve bias die zorgt voor robuustheid en het leren van onderliggende causale structuren.
• Ontwerp van Optimizers: Het ontwikkelen van nieuwe optimizers mag niet alleen focussen op convergentiesnelheid. De inductieve bias die een optimizer introduceert (bijv. via de manier waarop hij het verlieslandschap doorkruist) is fundamenteel voor het gedrag van het model.
• Toekomstige Richting: De auteurs pleiten voor optimizers die de voordelen van Muon (het vermijden van zadelpunten) combineren met de "simplicity bias" van SGD. Een ideale optimizer zou de trajecten van GD kunnen volgen, maar sneller door de zadelpunten heen kunnen breken zonder de sequentiële leerstructuur te verliezen.

Samenvattend: Muon is een krachtige tool voor snelheid, maar in complexe taken waar het vinden van gedeelde, onderliggende structuren of het vermijden van spurious correlaties cruciaal is, kan de traditionele SGD (vanwege zijn sequentiële, "simplicity-bias" leertraject) superieure resultaten opleveren.