Technische Samenvatting: Spectrale Flattening Is Alles Wat Muon Noodzaakt

Probleemstelling

Het trainen van moderne neurale netwerken houdt in dat geoptimaliseerd wordt over sterk niet-convexe landschappen die worden gekenmerkt door vlakke gebieden, scherpe richtingen en zadelpunten. De keuze van de optimizer is kritiek voor stabiliteit en efficiëntie. Hoewel coördinaatsgewijze adaptieve methoden zoals Adam en AdamW de robuustheid hebben verbeterd, lossen ze de geometrische uitdagingen die worden gesteld door matrixwaardige parameters niet volledig op. Onlangs is Muon naar voren gekomen als een hoogpresterend alternatief dat de momentumbuffer orthogonaliseert voorafgaand aan updates, wat groere leersnelheden en snellere convergentie mogelijk maakt. De theoretische mechanismen die uitleggen waarom Muon standaard Gradient Descent (SGD) overtreft in termen van maximaal stabiele stapgroottes en convergentiesnelheden, blijven echter onderbelicht. Specifiek heeft bestaande theorie niet kwantitatief uitgelegd hoe orthogonalisatie de bottleneck van de leersnelheid of het effectieve conditiegetal van het optimalisatieprobleem verandert.

Methodologie

Het artikel biedt een theoretische analyse van Muon door de update-regel ervan te kaderen als een vorm van spectrale flattening. De auteurs ontwikkelen twee primaire theoretische resultaten en valideren deze experimenteel:

1. Theoretisch Kader:

   • Muon Definitie: Muon werkt de gewichten $W$ bij door de gradient $G_t$ te berekenen en vervolgens Newton–Schulz-iteraties toe te passen om de polaire factor te benaderen $O_t \approx (G_t G_t^\top)^{-1/2} G_t = UV^\top$. Deze bewerking behoudt de singuliere vectoren van de gradient terwijl de singuliere waarden worden vervangen door enen.
   • Hessian Benadering: De analyse maakt gebruik van een Kronecker-gefactoreerd krommemodel (K-FAC), waarbij de Hessian wordt benaderd als $H \approx X^\top X \otimes G_t G_t^\top$, waarbij $X$ de input-activaties voorstelt en $G_t$ de gradient.
   • Leersnelheidsanalyse: De auteurs leiden exacte drempels voor één-staps afname af voor zowel SGD als Muon. Ze vergelijken de maximaal stabiele leersnelheid ($\eta_{max}$) voor beide optimalizers, en tonen aan dat SGD wordt geblokkeerd door de grootste singuliere waarde van de gradient ($\sigma_{max}$), terwijl de stabiliteit van Muon wordt bepaald door de gemiddelde singuliere waarde ($\frac{1}{m}\sum \sigma_i$).
   • Convergentieanalyse: Muon wordt herschreven als een gepreconditioneerde gradientmethode met preconditie $P = I_n \otimes (G_t G_t^\top)^{-1/2}$. Onder Polyak–Łojasiewicz (PL)-voorwaarden en relatieve gladheid analyseren de auteurs de convergentiefactor. Ze bewijzen dat de effectieve conditieverhouding voor Muon ($\tilde{\alpha}/\tilde{\beta}$) strikt beter is dan die van SGD ($\alpha/\beta$) wanneer de gradientcovariantie anisotroop is, specifiek met een verbetering die gerelateerd is aan de spectrale spreiding van de gradient.

2. Experimenteel Ontwerp:

   • Opstelling: Experimenten werden uitgevoerd op CIFAR-10 met CifarNet. Om het effect van de optimizer op matrixwaardige parameters te isoleren, werd een dual-optimizer strategie toegepast: convolutiekernen werden bijgewerkt door de doel-optimizer (Muon of SGD), terwijl biases en andere parameters een vaste referentie-optimizer gebruikten.
   • Stabiliteitstests: Leersnelheidsverkenningen werden uitgevoerd zonder Batch Normalization (BN) om de intrinsieke stabiliteit waar te nemen.
   • Convergentietests: Beide optimalizers werden uitgevoerd bij identieke leersnelheden (en afzonderlijk bij hun respectievelijke beste snelheden) om de convergentiesnelheid en per-stap convergentieverhoudingen ($r_t$) te meten.
   • Normalisatieprincipe: Een nieuw normalisatieschema, FrobNorm (normaliseren door de Frobenius-norm), werd voorgesteld en getest om de theoretische bewering te verifiëren dat het verkleinen van de spectrale schaal van inputs ($\lambda_{max}(X^\top X)$) hogere leersnelheden mogelijk maakt.

Belangrijkste Bijdragen

1. Afleiding van Maximaal Leersnelheid: Het artikel bewijst dat de maximaal stabiele stapgrootte van Muon schaalt met de gemiddelde singuliere waarde van de gradient in plaats van de grootste.

   • Voor SGD: $\eta_{max} \propto \frac{2}{\lambda_{max}(H)} \approx \frac{2}{\lambda_{max}(X^\top X) \sigma_{max}(G)}$.
   • Voor Muon: $\eta_{max} \propto \frac{2}{\lambda_{max}(H)} \frac{\sum \sigma_i}{m} \approx \frac{2}{\lambda_{max}(X^\top X) \sigma_{max}(G)} \frac{\sum \sigma_i}{m \sigma_{max}(G)}$.
   • Dit toont aan dat spectrale flattening voorkomt dat één dominante singuliere richting de stapgroottebeperking dicteert, waardoor Muon aanzienlijk grotere leersnelheden kan tolereren.

2. Versnelling van Convergentiesnelheid: De auteurs tonen aan dat Muon fungeert als een preconditie die de effectieve conditieverhouding verbetert.

   • Onder K-FAC is de convergentiefactor voor Muon $1 - \frac{\tilde{\alpha}}{\tilde{\beta}}$, waarbij $\frac{\tilde{\alpha}}{\tilde{\beta}} = \frac{\alpha}{\beta} \sqrt{\frac{\lambda_{min}(G_t G_t^\top)}{\lambda_{max}(G_t G_t^\top)}}$.
   • Aangezien $\frac{\lambda_{min}}{\lambda_{max}} < 1$, is de term $\sqrt{\dots}$ kleiner dan 1, maar de afleiding toont aan dat $\frac{\tilde{\alpha}}{\tilde{\beta}} > \frac{\alpha}{\beta}$ in de context van het gepreconditioneerde systeem, wat leidt tot een strikt snellere lineaire convergentiesnelheid wanneer de gradientcovariantie slecht geconditioneerd is.

3. Nieuw Normalisatieprincipe: De analyse onthult dat de maximaal stabiele leersnelheid omgekeerd evenredig is met $\lambda_{max}(X^\top X)$. De auteurs stellen voor dat normalisatielagen erop gericht moeten zijn om deze spectrale schaal te verkleinen. Ze valideren dit met FrobNorm, en tonen aan dat het training mogelijk maakt bij hoge leersnelheden (bijv. 0,25) waar standaard training faalt, met name voor SGD.

Experimentele Resultaten

• Stabiliteit: Bij afwezigheid van Batch Normalization divergeerde SGD binnen de eerste paar iteraties bij leersnelheden waarbij Muon stabiel bleef en de loss bleef verlagen. Muon vertoonde gecontroleerde groei in parameter- en gradientnormen, terwijl SGD een catastrofale explosie vertoonde.
• Convergentiesnelheid: Bij identieke leersnelheden (bijv. $\eta=0,05$) bereikte Muon mijlpalen in nauwkeurigheid (70%–91%) enkele epochs eerder dan SGD. De empirische convergentieverhouding $r_t$ was consequent lager voor Muon, wat de theoretische voorspelling van een snellere lineaire snelheid bevestigt.
• Impact van Normalisatie: De toepassing van FrobNorm stelde beide optimalizers in staat om stabiel te trainen bij hoge leersnelheden (0,25 voor CifarNet, 0,2 voor een GPT-2-achtige Transformer), wat de hypothese bevestigt dat het beheersen van de spectrale schaal van lageninputs een algemeen principe is om grotere stapgroottes mogelijk te maken.

Betekenis en Claims

Het artikel claimt een gebaseerde, geometrische verklaring te bieden voor het empirische succes van Muon. Het stelt dat het voordeel van Muon niet louter een heuristiek is, maar een direct gevolg van spectrale flattening, wat:

1. De leersnelheidsbeperking ontkoppelt van de grootste singuliere waarde, waardoor grotere stapgroottes mogelijk zijn.
2. Fungeert als een eenzijdige preconditie die het effectieve conditiegetal van het optimalisatielandschap verbetert, wat leidt tot snellere convergentie.

De auteurs erkennen beperkingen, en merken op dat hun theoretische analyse deterministische full-batch instellingen en exacte polaire factoren veronderstelt, terwijl praktische implementaties gebruikmaken van stochastische gradienten, eindige Newton–Schulz-iteraties en momentum. Ze stellen dat het uitbreiden van deze resultaten naar stochastische instellingen en grootschaligere architecturen een belangrijke richting blijft voor toekomstig werk. Het artikel concludeert dat deze resultaten helpen het empirische vermogen van Muon om te zetten in een preciezere theoretische understanding, wat potentieel het ontwerp van toekomstige optimalizers en normalisatiestrategieën kan sturen.