How Much Orthogonalization Does Muon Need?

Technische Samenvatting: Hoeveel orthogonalisatie heeft Muon nodig?

Probleemstelling
De Muon-optimizer verbetert de training van neurale netwerken door ill-conditioned momentum-updates te vervangen door benaderende semi-orthogonale updates. Terwijl standaard numerieke lineaire algebra streeft naar een hoge nauwkeurigheid in de polaire decompositie (het mappen van alle singuliere waarden naar 1), suggereert de Muon-context dat exacte nauwkeurigheid niet noodzakelijk is. De centrale vraag die dit werk adresseert is: Hoeveel nauwkeurigheid in orthogonalisatie heeft Muon daadwerkelijk nodig om de kwaliteit van de training te behouden? Specifiek onderzoekt het paper of de computationele kosten van hogere-orde polynoom-iteraties (gebruikt in bestaande methoden zoals Muon-Jordan en Polar Express) gerechtvaardigd zijn, of dat een goedkopere, lager-orde benadering vergelijkbare trainingsresultaten kan behalen.

Methodologie
De auteurs stellen een ontspannen kubische Newton–Schulz-schema voor (termed cubic5) die specif gezien is ontworpen voor de bfloat16 low-precision singular value band van Muon.

1. Ontspannen Targets: In tegenstelling tot klassieke polaire decompositie die streeft naar machineprecisie, richt deze methode zich op een ontspannen band van $[0,7, 1,3]$ voor singuliere waarden. Het erkent dat in bfloat16 singuliere waarden onder een effectieve ondergrens ($l_0 = 7 \times 10^{-3}$) niet praktisch betekenisvol zijn.
2. Adaptieve Coëfficiëntafleiding: De auteurs leiden een adaptieve kubische polynoom $f_t(x) = a_t x + b_t x^3$ af voor elke iteratie. De coëfficiënten worden berekend op basis van de huidige worst-case lagere singuliere waarde-ondergrens ($l_t$) en een vaste bovenpiek ($u=1,3$). Dit zorgt ervoor dat de polynoom het beschermde interval $[l_t, r_t]$ mapt naar de doelband, terwijl de eindpunten worden gebalanceerd om de nieuwe ondergrens te maximaliseren.
3. Computationele Efficiëntie:
   • Cubic5: Een enkele kubische stap vereist twee dominante matrixvermenigvuldigingen (het vormen van de Gram-matrix en het terugvermenigvuldigen ermee). Vijf stappen vereisen 10 dominante matrixvermenigvuldigingen.
   • Baseline (Muon-Jordan/Polar Express): Een quintische stap vereist drie dominante matrixvermenigvuldigingen. Vijf stappen vereisen 15 dominante matrixvermenigvuldigingen.
   • De methode maakt gebruik van de "transpose trick" voor hoge matrices om operaties op de kleinere Gram-matrix te minimaliseren.

Belangrijkste Bijdragen

• Afleiding van een Ontspannen Kubisch Schema: Het paper leidt een vijfstaps kubisch Newton–Schulz-schema af met coëfficiënten die zijn aangepast aan de bfloat16 effectieve ondergrens, waarmee het doel spectrale band bereikt in vijf iteraties.
• Karakterisering van de Cost-Quality Frontier: Het werk stelt vast dat de kwaliteit van de training niet monotoon wordt bepa wordt door de nauwkeurigheid van de polaire decompositie. Het demonstreert dat een goedkopere, lagere-orde kubische benadering in staat is om prestatie-pariteit te bereiken met hogere-kosten quintische methoden en zelfs exacte SVD in specifieke regimes.
• Empirische Validatie over Verschillende Schalen: De auteurs valideren de methode via:
  • Synthetische Diagnostiek: Bevestiging dat de scalaire polynoomtransfer standhoudt in low-precision matrix-arithmetica.
  • NanoGPT Ablaties: Vergelijking van cubic5, getrunceerde Polar Express, getrunceerde Muon-Jordan en expliciete FP32 SVD op GPT-2 Small.
  • Grootschalige Training: Training van hybride MoE/Mamba-modellen (Nemotron-3-Nano) variërend van 1B tot 4B parameters met behulp van Megatron-LM.

Resultaten

• Synthetische Benchmarks: De cubic5-methode mapt succesvol singuliere waarden in het beschermde bereik $[7 \times 10^{-3}, 1]$ naar de doelband zonder numerieke explosie. Microbenchmarks op NVIDIA H200 GPU's laten zien dat cubic5 de kosten van de orthogonalisatie-subroutine vermindert ten opzichte van Polar Express, met versnellingen variërend van 1,16x tot 1,52x, afhankelijk van de matrixvorm (rechthoekig vs. vierkant).
• GPT-2 Small Ablaties:
  • Vijfstaps cubic5 bereikt een uiteindelijke validatieverlies van 3,2811.
  • Vijfstaps Polar Express bereikt 3,2798.
  • Vijfstaps Muon-Jordan bereikt 3,2822.
  • Expliciete FP32 SVD bereikt 3,2809.
  • Observatie: De kloof tussen de sterkste varianten is minder dan 0,003 in validatieverlies. Opvallend genoeg presteert de exacte SVD niet beter dan de benaderende polynoom-methoden, en het toevoegen van een vijfde stap aan Muon-Jordan degradeerde de prestaties licht vergeleken met vier stappen.
• Hybride MoE/Mamba Schaling (1B–4B):
  • Op 1B en 2B modellen presteerde cubic5 licht beter dan Muon-Jordan (lager validatieverlies met $\approx 10^{-3}$).
  • Op het 4B model was cubic5 iets hoger ($\approx 10^{-3}$), maar de trainingspaden waren vrijwel ononderscheidbaar na de initiële fase.
  • Het relatieve verschil in trainingsverlies nam snel af en fluctueerde rond nul gedurende de rest van het trainingsbudget.

Betekenis en Claims
Het paper claimt bescheiden dat cubic5 een praktische, goedkope variant is voor Muon-orthogonalisatie, in plaats van een universeel superieure update-regel.

• Optimizer vs. Decompositie Nauwkeurigheid: De resultaten ondersteunen de hypothese dat Muon profiteert van een "spectrale herformulerings"-effect in plaats van strikte polaire decompositie-nauwkeurigheid. De exacte singuliere waarde-mapping die de training optimaliseert, hangt af van de trainingsfase en de learning rate, niet alleen van wiskundige convergentie naar 1.
• Praktisch Nut: cubic5 biedt een principieel punt op de cost-quality frontier. Het vermindert het aantal dominante matrixvermenigvuldigingen met ongeveer een derde (10 vs. 15) terwijl het empirische kwaliteit van de training gelijk houdt aan de standaard Muon-Jordan quintische update in de geteste regimes.
• Beperkingen: De auteurs geven expliciet aan dat deze resultaten training-kwaliteit pariteit vaststellen, geen end-to-end throughput versnelling. De werkelijke wall-clock versnelling hangt af van hardware kernels, batching en het deel van de tijd dat de orthogonalisatiestap de volledige optimizer-loop inneemt. Het paper beweert niet dat cubic5 een beter polaire decompositie algoritme is dan Polar Express in een wiskundige zin, maar eerder een efficiëntere spectrale transformatie voor de specifieke context van Muon-training.