Mousse: Rectifying the Geometry of Muon with Curvature-Aware Preconditioning

Dit paper introduceert Mousse, een nieuwe optimizer die de stabiliteit van Muon combineert met de geometrische aanpassing van Shampoo door spectrale updates uit te voeren in een witgemaakte coördinatenruimte, wat leidt tot een aanzienlijke versnelling van het trainingsproces voor taalmodellen zonder noemenswaardige rekenkosten.

De kern

Stel je voor dat je een enorme, complexe berglandschap moet doorkruisen om de laagste vallei (de beste oplossing) te vinden. In de wereld van kunstmatige intelligentie is dit landschap het "verlieslandschap" van een groot taalmodel. De uitdaging is dat dit landschap niet gelijkmatig is: sommige hellingen zijn steil en gevaarlijk, andere zijn heel plat en saai, en weer andere zijn vol met gaten en kuilen.

Dit artikel introduceert Mousse, een nieuwe "kompas" (een optimalisator) die veel slimmer is dan de huidige standaard, genaamd Muon.

Hier is de uitleg in simpele taal, met een paar creatieve vergelijkingen:

1. Het Probleem: De "Gelijke-Behandelings"-Fout

Stel je voor dat je een groep wandelaars hebt die allemaal precies hetzelfde stappenplan volgen, ongeacht het terrein.

• Muon (de oude methode) doet dit. Het zegt: "We maken een stap, en die stap mag overal even groot zijn, maar we mogen niet te steil klimmen." Het behandelt elke richting op de berg alsof ze even gevaarlijk zijn.
• Het probleem: In werkelijkheid is dat niet zo. Sommige richtingen zijn als een gladde ijsbaan (zeer gevoelig, je valt snel), en andere zijn als een zanderig vlak (je moet hard werken om vooruit te komen). Als Muon op de ijsbaan dezelfde stapgrootte neemt als op het zand, glijdt je uit of loop je vast. Het is te "egalitair" (gelijkheidszuchtig) voor een ongelijk landschap.

2. De Oplossing: Mousse (De Slimme Gids)

Mousse (wat staat voor Muon Optimization Utilizing Shampoo's Structural Estimation) is als een wandelaar met een 3D-kaart en een GPS die het terrein in real-time scant.

In plaats van gewoon te stappen, doet Mousse drie dingen:

1. Het landschap "witten" (Sphering): Mousse kijkt eerst naar de vorm van de berg. Het rekent het landschap zo uit dat de steile hellingen en de vlakke dalen even "normaal" lijken. Het is alsof je een vervormde foto van de berg neemt en die uitrekt tot een perfect vlakke kaart.
2. De stap nemen: Op die nieuwe, uitgebalanceerde kaart neemt Mousse de perfecte stap (net zoals Muon dat doet, maar dan op de juiste manier).
3. Terugrekenen: Daarna zet Mousse de stap weer om naar de echte, vervormde berg.

De metafoor:
Stel je voor dat je een trui moet strijken.

• Muon strijkt de trui alsof hij perfect plat ligt. Als de stof kreukelig is, krijg je geen goed resultaat.
• Mousse pakt eerst een speciale strijkplaat die precies de vorm van de kreukels volgt (de kromming van de stof). Dan strijkt hij eroverheen. Het resultaat is een perfect gladde trui, veel sneller en beter dan Muon.

3. Waarom is dit zo goed?

De onderzoekers hebben dit getest op modellen van 160 miljoen tot 800 miljoen "woorden" (parameters). De resultaten zijn indrukwekkend:

• Sneller: Mousse heeft ongeveer 12% minder stappen nodig om hetzelfde resultaat te bereiken als Muon. Dat is alsof je een reis van 100 kilometer in 88 kilometer doet.
• Beter: De wandelaars komen uiteindelijk in een diepere vallei (een lagere foutwaarde).
• Efficiënt: Het kost bijna evenveel rekenkracht als Muon. Het is niet zwaar en traag, zoals andere geavanceerde methoden (zoals SOAP) die vaak te veel geheugen nodig hebben.

4. De Geheime Ingrediënten (De "Trucs")

Om dit te laten werken zonder dat de wandelaars uit de bocht vliegen, gebruiken ze twee slimme trucjes:

• Trace Normalization: Dit zorgt ervoor dat de kaart niet te groot of te klein wordt. Het houdt de schaal in balans, zodat de gids niet in paniek raakt als de berg plotseling heel hoog lijkt.
• Spectral Tempering: Dit is als een "rem" op de remmen. Soms is de kaart zo scherp dat je te hard zou gaan. Mousse dempt deze scherpe hoeken een beetje, zodat je stabiel blijft rennen zonder te vallen.

Conclusie

Kortom: Mousse is de volgende generatie wandelgids voor AI. Het combineert de snelheid en stabiliteit van de oude Muon-methode met de slimme landschapsanalyse van geavanceerde wiskundige methoden. Het zorgt ervoor dat AI-modellen sneller leren, minder energie verbruiken en betere resultaten leveren, zonder dat het de computer te veel kost.

Het is een bewijs dat je niet altijd harder hoeft te werken, maar dat je soms gewoon een beter kompas nodig hebt.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "Mousse: Rectifying the Geometry of Muon with Curvature-Aware Preconditioning" in het Nederlands.

Titel: Mousse: Rectificatie van de Geometrie van Muon met Krommingsbewuste Preconditionering

Datum: 11 maart 2026
Auteurs: Yechen Zhang et al. (Shanghai Jiao Tong University, Shanghai AI Laboratory, Fudan University)

---

1. Het Probleem

Recente vooruitgang in spectrale optimalisatie, met name de Muon-optimizer, heeft aangetoond dat het beperken van update-stappen tot de Stiefel-maand (via Newton-Schulz-iteraties) het trainen van grote taalmodellen kan versnellen en de generalisatie kan verbeteren. Muon werkt echter op basis van een cruciaal geometrisch aanneming: het veronderstelt een isotroop optimalisatielandschap. Dit betekent dat het een uniforme spectrale update-norm oplegt over alle eigenrichtingen.

De auteurs betogen dat deze "egalitaire" beperking suboptimaal is voor Deep Neural Networks (DNN's), waar het krommingsspectrum (curvature spectrum) bekend staat om zijn hoog zwaarstaartgedrag (heavy-tailed) en slecht geconditioneerdheid (ill-conditioned). In dergelijke landschappen riskeert Muon:

• Instabiliteiten te versterken in richtingen met hoge kromming.
• Noodzakelijke vooruitgang te beperken in vlakke richtingen.

De kernuitdaging is om Muon's rigide isotrope beperking te verenigen met de sterk anisotrope kromming van neurale netwerken zonder de rekenkosten van tweede-orde methoden (zoals Shampoo) volledig te hoeven dragen.

2. Methodologie: De Mousse-Optimizer

Mousse (Muon Optimization Utilizing Shampoo's Structural Estimation) is een nieuwe optimizer die de structurele stabiliteit van spectrale methoden combineert met de geometrische aanpasbaarheid van tweede-orde preconditionering.

Kernidee:
In plaats van Newton-Schulz-orthogonalisatie direct toe te passen op de momentum-matrix in de oorspronkelijke coördinaten, voert Mousse de optimalisatie uit in een ge-wit (whitened) coördinatenstelsel. Dit wordt bereikt door de gradiënt eerst te preconditioneren met statistieken van de Kronecker-factoren (afgeleid van de Shampoo-optimizer).

Wiskundige Formulering:

1. Geometrische Rectificatie: Mousse reformuleert het optimalisatieprobleem als een spectrale steilste afdaling (spectral steepest descent) onder een anisotrope vertrouwensregio (trust region).
2. Whitening: De Hessian-benadering wordt gesplitst in Kronecker-factoren $L$ (rijen) en $R$ (kolommen). De gradiënt $G$ wordt getransformeerd naar een witte ruimte via:
   $$\tilde{G} = L^{-1/4} G R^{-1/4}$$
3. Spectrale Projectie: De Newton-Schulz-iteratie (msign) wordt toegepast op deze getransformeerde gradiënt $\tilde{G}$ om de update-richting te bepalen.
4. Terugprojectie: De update wordt teruggeprojecteerd naar de oorspronkelijke ruimte:
   $$\Delta W = -L^{-1/4} \cdot \text{msign}(L^{-1/4} G R^{-1/4}) \cdot R^{-1/4}$$

Dit zorgt ervoor dat de spectrale beperking (die de Stiefel-maand respecteert) wordt toegepast op een landschap dat lokaal "bol" (sferisch) is gemaakt door de krommingsinformatie.

Technische Innovaties voor Stabiliteit:
Om de numerieke stabiliteit van deze tweede-orde spectrale methode te garanderen, introduceert het paper twee kritieke technieken:

• Trace Normalization: Normalisatie van de covariantiematrices ($L$ en $R$) zodat het gemiddelde eigenwaarde 1 is. Dit zorgt voor consistente demping over verschillende lagen.
• Spectral Tempering: Het gebruik van een milder exponent ($\alpha = 0.125$ in plaats van de standaard $0.25$) bij het berekenen van de negatieve macht van de eigenwaarden. Dit voorkomt overmatige versterking van gradiënten in vlakke richtingen (waar eigenwaarden klein zijn).
• Gradient Grafting: Het koppelen van de update-grootte aan een stabiele methode (zoals AdamW) om te voorkomen dat de update-norm tijdens de training afneemt.

3. Belangrijkste Bijdragen

1. Unificatie van Geometrie: Mousse biedt een theoretisch onderbouwde unificatie van spectrale optimalisatie en tweede-orde preconditionering. Het lost het fundamentele probleem op dat spectrale methoden vaak de kromming van het landschap negeren.
2. Robuuste Ingenieursinzichten: Het paper biedt praktische richtlijnen voor het stabiliseren van tweede-orde spectrale optimalisatie, specifiek door Trace Normalization en Spectral Tempering.
3. Pareto-optimale Efficiëntie: Mousse bereikt een nieuw evenwicht tussen prestaties en rekenkosten, waarbij het de voordelen van Shampoo haalt zonder de zware geheugen- en rekentijdstraffen.

4. Resultaten

Experimenten zijn uitgevoerd op taalmodellen variërend van 160M tot 800M parameters getraind op de FineWeb-dataset.

• Verbeterde Convergentie: Mousse overtreft consistent de standaard Muon-optimizer. Op het 800M-model wordt een reductie van ongeveer 12% in het aantal trainingsstappen bereikt om een vergelijkbaar verliesniveau te bereiken.
• Verliesreductie: Mousse bereikt een lager eindverificatieverlies (validation loss) dan Muon, SOAP en AdamW over alle modelgroottes.
• Rekenkosten:
  • Tijd: De wall-clock tijd-overhead is verwaarloosbaar (ongeveer 3% ten opzichte van Muon).
  • Geheugen: Mousse is aanzienlijk geheugenefficiënter dan SOAP (ongeveer 88% van het geheugengebruik van SOAP) en vergelijkbaar met Muon, omdat het geen tweede-moment buffer (variance buffer) nodig heeft.
• Schaalbaarheid: De prestatieverbetering is robuust over verschillende modelgroottes en leerplans (learning rates).

5. Betekenis en Impact

Dit werk markeert een belangrijke stap in de evolutie van optimizers voor grote taalmodellen (LLMs).

• Oplossing voor het "Isotrope Dilemma": Het toont aan dat spectrale methoden (zoals Muon) niet inherent superieur zijn, maar dat hun prestaties drastisch verbeteren wanneer ze worden gecombineerd met krommingsbewuste preconditionering.
• Efficiëntie voor Pre-training: Door de 12% reductie in trainingsstappen met minimale extra kosten, biedt Mousse een directe route naar kostenbesparing en snellere ontwikkeling van grote AI-modellen.
• Toekomstperspectief: De methode opent de deur voor verdere integratie van geavanceerde preconditioneringstechnieken (zoals single-sided preconditioners) en mogelijke verbeteringen in fine-tuning-scenario's.

Kortom, Mousse corrigeert de geometrische tekortkomingen van Muon door de structuur van Shampoo te integreren, wat resulteert in een optimizer die zowel sneller convergeert als rekenkundig efficiënter is dan bestaande state-of-the-art methoden.