Delving into Muon and Beyond: Deep Analysis and Extensions

Technische Samenvatting: Een diepe analyse en uitbreiding van Muon en verder

Probleemstelling
De Muon-optimizer heeft recentelijk veel aandacht gekregen vanwege de empirische prestaties bij het trainen van grote taalmodellen, waarbij gebruik wordt gemaakt van matrix-niveau georthogonaliseerde updates. De onderliggende mechanismen van Muon en de relatie tot gevestigde adaptieve optimizers zoals Adam blijven echter onvoldoende begrepen. Bestaande literatuur verwart de effecten van Muon vaak met hulptechnieken (bijv. QK-Norm, QK-Clip) of mist een systematische vergelijking met standaard baselines. Bovendien is het onduidelijk of de spectrale orthogonalisatie van Muon een universeel superieure optimalisatiestrategie vertegenwoordigt of een specifieke vorm van stabilisatie, en of het verbeterd of gegeneraliseerd kan worden.

Methodologie
De auteurs stellen een verenigd spectraal kader voor om Muon en zijn varianten te analyseren. Ze definiëren Muon als het $p=0$ eindpunt van een familie van spectrale transformaties van de vorm $U\Sigma^p V^\top$, waarbij $U\Sigma V^\top$ de Singular Value Decomposition (SVD) is van de update-matrix. De studie onderzoekt tussenliggende exponenten $p \in \{1, 1/2, 1/4, 0\}$:

• $p=1$: Identiteit (standaard updates).
• $p=1/2, 1/4$: Partiële spectrale compressie.
• $p=0$: Volledige orthogonalisatie (Muon).

Deze transformaties worden toegepast op twee verschillende soorten inputs:

1. First-moment momentum ($O_{mom}$): Equivalent aan momentum SGD-inputs.
2. RMS-genormaliseerde updates ($O_{rms}$): Equivalent aan Adam-inputs, waarbij gradiënten worden genormaliseerd door het tweede moment.

Om efficiënte berekeningen van fractionele exponenten ($p=1/2, 1/4$) mogelijk te maken zonder expliciete SVD, ontwikkelen de auteurs een gekoppelde Newton-Schulz iteratie. Dit algoritme berekent simultaan de matrixwortel en de inverse ervan met enkel matrixvermenigvuldigingen, wat numerieke stabiliteit en schaalbaarheid garandeert.

De experimentele evaluatie wordt uitgevoerd op een GPT-2 model met 124 miljoen parameters getraind op OpenWebText. De studie hanteert een strikt gecontroleerde opzet:

• Ontkoppelde leersnelheden: Er worden aparte leersnelheden gebruikt voor matrix- en vectorparameters om het effect van de optimizer te isoleren.
• Geen hulptechnieken: Weight decay, QK-Norm en QK-Clip zijn uitgeschakeld om eerlijke vergelijkingen te waarborgen.
• Systematische tuning: Een tweestaps leersnelheid-zoektocht (grove logaritmische sweep gevolgd door fijne verfijning) wordt uitgevoerd voor alle acht optimizer-varianten (4 spectrale exponenten $\times$ 2 inputtypen).

Belangrijkste bijdragen

1. Verenigd Spectraal Kader: Het artikel introduceert een formeel kader $U\Sigma^p V^\top$ dat Muon positioneert als een specifiek geval ($p=0$) en stelt tussenliggende varianten ($p=1/2, 1/4$) voor om te interpoleren tussen standaard updates en volledige orthogonalisatie.
2. Efficiënt Algoritme: Een gekoppelde Newton-Schulz iteratie is ontwikkeld om fractionele spectrale updates ($p=1/2, 1/4$) efficiënt te berekenen, waardoor de computationele kosten van expliciete SVD worden vermeden.
3. Rigoureuze Empirische Vergelijking: De studie biedt een gecontroleerde vergelijking tussen Muon, Adam en hun spectrale varianten, waarbij spectrale effecten worden geïsoleerd van verwarrende factoren zoals weight decay of hulp-normalisatie.

Belangrijkste Resultaten
De experimenten leveren drie primaire bevindingen op:

1. Stabilisatie van First-Moment Updates: Op first-moment inputs (mSGD), biedt Muon ($p=0$) significante stabilisatie, wat veel hogere leersnelheden en grotere robuustheid mogelijk maakt vergeleken met standaard mSGD. Echter, partiële compressie ($p=1/4$) bereikt vaak betere uiteindelijke prestaties zodra de stabiliteit is gevestigd.
2. Beperde Winst op RMS-genormaliseerde Inputs: Wanneer toegepast op RMS-genormaliseerde inputs (Adam-stijl), levert spectrale compressie afnemende meeropbrengsten op. Volledige orthogonalisatie (AdamZ/Muon toegepast op $O_{rms}$) presteert niet beter dan standaard Adam en kan inferieur zijn aan partiële compressie. Dit suggereert dat elementgewijze RMS-normalisatie al veel van de anisotropie die spectrale compressie aanpakt, reeds mitigeert.
3. Superioriteit van Second-Moment Methoden: Over alle spectrale exponenten heen, presteren optimizers die gebruikmaken van second-moment normalisatie (RMS-gebaseerd) consistent beter dan hun first-moment (alleen momentum) tegenhangers in zowel stabiliteit als uiteindelijke loss.

Betekenis en Claims
Het paper concludeert dat Muon het beste begrepen kan worden als een effectieve vorm van spectrale normalisatie in plaats van een universeel superieure optimalisatiemethode.

• Stabilisatie versus Prestatie: De primaire bijdrage van Muon is het stabiliseren van first-moment updates door de magnitudes van singuliere waarden te verwerpen, wat voorkomt dat richtingen met grote singuliere waarden domineren. Deze "afvlakking" is echter een grof mechanisme; het verwijdert magnitude-gebaseerde filtering, wat potentieel ruisgevoelige gradiëntrichtingen (kleine singuliere waarden) versterkt en onderstappen veroorzaakt op informatieve dominante modi.
• Contextuele Nuttigheid: Hoewel spectrale compressie echte voordelen biedt voor stabilisatie, vooral wanneer de input-update niet al genormaliseerd is, is volledige orthogonalisatie ($p=0$) geen vervanging voor moderne RMS-genormaliseerde second-moment optimizers zoals Adam. De auteurs stellen dat de gerapporteerde voordelen van Muon in eerder werk mogelijk toe te schrijven zijn aan de specifieke combinatie van technieken (bijv. QK-Norm) in plaats van aan de orthogonalisatie zelf.

Het werk benadrukt dat hoewel spectrale transformaties krachtige instrumenten zijn voor het controleren van anisotropie, ze inherent niet de voordelen van adaptieve second-moment schaling in grootschalige trainingsregimes overtreffen.