Technische Samenvatting: Muonp – Muon met Fractionele Spectrale Machten

Probleemstelling

De Muon-optimizer is een krachtig instrument geworden voor het trainen van diepe neurale netwerken, in het bijzonder grote taalmodellen (LLM's). Het werkt door de gradiëntmatrix $G = USV^\top$ te vervangen door zijn polaire factor $UV^\top$, wat effectief de spectrale verdeling van de singuliere waarden afvlakt. Deze "spectrumafvlakking" bevordert gebalanceerd leren over singuliere richtingen en voorkomt dat dominante modi de trainingsupdates monopoliseren.

De auteurs identificeren echter een kritieke beperking in de aanpak van Muon: volledige afvlakking verwerpt alle informatie over de singuliere waarden. Hoewel deze uniformiteit gunstig is voor brede feature-ontdekking tijdens pretraining, kan het suboptimaal zijn voor adaptatietaken (zoals finetuning), waarbij de grootte van de singuliere waarden aangeeft welke residuele richtingen het meest waard zijn om te veranderen. De centrale vraag die dit werk adresseert is: Kunnen we de spectrale voordelen van Muon op matrixniveau behouden terwijl we de magnitudestructuur behouden die noodzakelijk is voor effectieve adaptatie?

Methodologie

Theoretische Formulering: Muonp

De auteurs introduceren Muonp, een generalisatie van de Muon-optimizer die gewichten bijwerkt met behulp van fractionele spectrale machten:
$$\text{Update} \propto US^pV^\top$$
waarbij $p \in (0, 1)$ een rationaal getal is.

• Geometrische Interpretatie: De auteurs bewijzen (Theorema 2.1) dat deze update-regel overeenkomt met de steilste-afdaling-richting onder de Schatten $q$-norm, waarbij $q = 1 + 1/p$.
• Interpolatie:
  • $p = 0$: Komt overeen met standaard Muon ($UV^\top$), wat uniformiteit maximaliseert.
  • $p = 1$: Komt overeen met standaard gradiëntafdaling ($USV^\top = G$).
  • $0 < p < 1$: Interpoleert tussen beide, waarbij gedeeltelijke magnitudainformatie behouden blijft.

Algoritmische Uitdaging en Oplossing

Het direct berekenen van $US^pV^\top$ vereist Singular Value Decomposition (SVD), wat computationeel onhaalbaar is voor grote modellen. De auteurs pakken de uitdaging aan om deze update te benaderen zonder expliciete SVD:

1. Onmogelijkheid van Univariante Iteratie: De auteurs bewijzen (Lemma 2.2) dat geen enkele vaste univariante polynoomiteratie $S_{n+1} = f(S_n)$ kan convergeren naar $S^p$ voor willekeurige rationale $p$. In tegen tegenstelling tot de polaire factor ($p=0$), vereisen fractionele machten dat het algoritme de oorspronkelijke matrix "onthoudt".
2. Bivariante Polynoom-recurrente: Om dit op te lossen, leiden de auteurs laaggradige oneven bivariante recurrente regels af. Het algoritme berekent iteratief matrices $G_n = US_nV^\top$ waarbij de singuliere waarden evolueren via $S_{n+1} = f(S_n, S)$.
   • Theorema 2.3 demonstreert dat voor elke oneven bivariante polynoom $f(x, y)$, de term $Uf(S_n, S)V^\top$ uitsluitend berekend kan worden met matrixvermenigvuldigingen en transposities, waardoor de computationele structuur van Muon behouden blijft.
   • Theorema 2.4 stelt de existentie vast van dergelijke polynomen die convergeren naar $S^p$ voor elke rationale $p \in (0, 1)$.
3. Expliciete Constructie: Het artikel geeft expliciete formules voor deze polynomen (bijv. $f(x, y) = x + c(y - x^3)$ voor $p=1/3$). De coëfficiënten worden gekozen om te garanderen dat $S^p$ een aantrekkend vast punt is.

Implementatiedetails

• Complexiteit: De computationele complexiteit blijft $O(\omega k N)$, waarbij $\omega$ de complexiteit van matrixvermenigvuldiging is, $k$ het aantal vermenigvuldigingen per stap (afhankelijk van de graad van de polynoom), en $N$ het aantal iteraties. Muonp verschilt van Muon alleen door een constante factor in het aantal matrixvermenigvuldigingen.
• Stabiliteit: De methode is numeriek stabiel wanneer deze in bfloat16 wordt geïmplementeerd, wat het geschikt maakt voor moderne hardware.
• Curriculum Compatibiliteit: Omdat Muon ($p=0$) een speciaal geval is van Muonp, vereist het wisselen tussen hen tijdens de training (bijv. pretraining met Muon en finetuning met Muonp) geen reset van de status, enkel een verandering in de polynoombenadering.

Belangrijkste Bijdragen

1. Muonp Optimizer: Introductie van een gegeneraliseerde optimizer die $US^pV^\top$ berekent voor elke rationale $p \in (0, 1)$ zonder expliciete SVD.
2. Theoretische Limieten: Bewijs dat fractionele spectrale machten niet berekend kunnen worden via vaste univariante polynoomiteraties, wat een nieuwe bivariante aanpak noodzakelijk maakt.
3. Constructief Algoritme: Afleiding van laaggradige oneven bivariante recurrente regels die fractionele machten benaderen terwijl de matrixvermenigvuldigingsstructuur van Muon behouden blijft.
4. Empirische Validatie: Demonstratie dat Muonp significant beter presteert dan Muon en AdamW in finetuning-scenario's voor wiskundige redenering, coderen en downstream taken.
5. Spectrale Geometrie Analyse: Inzicht dat de optimale spectrale geometrie regime-afhankelijk is: uniforme updates ($p=0$) zijn gunstig voor pretraining, terwijl gedeeltelijke preservatie ($0 < p < 1$) gunstig is voor finetuning.

Experimentele Resultaten

Finetuning Prestaties

Muonp vertoont consistente verbeteringen ten opzichte van Muon en AdamW bij het finetunen van Llama3-1B en Llama3-3B modellen:

• Wiskunde & Code Redeneren: Op GSM8K, Numina en OpenCodeInstruct verbetert Muonp de pass@8 nauwkeurigheid en vermindert het de loss vergeleken met Muon.
• Perplexity: Muonp bereikt een lagere validatie-perplexity op de Fineweb en Pile of Law datasets.
• Downstream Taken: Muonp levert een hogere nauwkeurigheid op een reeks redeneerbenchmarks (PIQA, Winogrande, LAMBADA, etc.) vergeleken met zowel Muon als AdamW.
• Optimizer Mismatch: Zelfs wanneer er met Muon is gepretrained (overeenkomstige optimizer), presteert Muonp beter dan Muon tijdens de finetuning, wat suggereert dat de winst fundamenteel is aan de spectrale geometrie in plaats van slechts optimizer-consistentie.

Pretraining vs. Finetuning Dichotomie

Het artikel onthult een duidelijke dichotomie in prestaties op basis van de trainingsfase:

• Pretraining: Muon ($p=0$) presteert beter dan Muonp. De auteurs stellen dat pretraining profiteert van "isotropisering" om brede features te ontdekken en te voorkomen dat dominante richtingen de updates monopoliseren.
• Finetuning: Muonp ($0 < p < 1$) presteert beter dan Muon. In dit regime heeft het model een lage intrinsieke dimensie, en helpt het behouden van de singuliere waarde-informatie om de adaptatie te focussen op de meest relevante residuele richtingen.
• Curriculum Experimenten: Het wisselen van Muon naar Muonp in de laatste stappen van de pretraining (een curriculum) levert onmiddellijke reducties in de loss op, wat de hypothese valideert dat het verhogen van $p$ de effectieve leersnelheid voor kleinere singuliere richtingen aanpast.

Runtime

Muonp behoudt een runtime die vergelijkbaar is met die van Muon. Bijvoorbeeld, met $p=1/3$ (cubische recurrente), is Muonp iets sneller dan Muon (dat in sommige implementaties een quintische recurrente gebruikt voor $p=0$), met staptijden rond de 100–120ms op H100 GPU's.

Betekenis en Claims

Het artikel positioneert Muonp als een geprincieerde generalisatie van de Muon-optimizer. De primaire betekenis ligt in:

1. Overbruggen van de Kloof: Het biedt een systematische manier om te interpoleren tussen de uniformiteit van Muon en de magnitudesensitiviteit van gradiëntafdaling, waardoor de trade-off tussen "spectrumafvlakking" en "magnitude-retentie" wordt opgelost.
2. Regime-Afhankelijke Optimalisatie: Het levert bewijs dat de "beste" spectrale geometrie niet universeel is; pretraining vereist volledige afvlakking, terwijl finetuning gedeeltelijke afvlakking vereist.
3. Praktische Efficiëntie: Het bereikt deze theoretische voordelen zonder de kosten van SVD op te lopen, waarbij de hardware-vriendelijke, enkel op matrixvermenigvuldiging gebaseerde structuur behouden blijft die essentieel is voor schalen naar miljarden parameters.
4. Toekomstige Richtingen: De auteurs suggereren dat het automatisch selecteren van de optimale $p$ voor specifieke taken en het integreren van Muonp in pretraining-curricula veelbelovende paden zijn voor toekomstig onderzoek.

Het werk concludeert dat het behouden van het singuliere spectrum aanzienlijke winsten kan brengen in adaptatiescenario's, en dat Muonp een computationeel efficiënte, geprincieerde methode biedt om dit te bereiken.