The Spectral Dynamics and Noise Geometry of Muon

Technische Samenvatting: De Spectrale Dynamiek en Ruisgeometrie van Muon

Probleemstelling
Het artikel onderzoekt de impliciete regularisatiebias geïnduceerd door de Muon-optimizer, die gewichten bijwerkt met behulp van de polaire factor van de gradiëntmatrix in plaats van de gradiënt zelf. Gegeven een gradiënt $G = U\Sigma V^\top$, werkt Muon bij via $UV^\top$, waardoor de relatieve grootte van de singuliere waarden wordt genegeerd terwijl de actieve singuliere richtingen behouden blijven. Hoewel Muon empirisch succes heeft getoond bij transformer-pretraining, blijft het theoretische selectiemechanisme onduidelijk. Twee concurrerende hypothesen worden behandeld:

1. Is Muon louter een vorm van matrix-gradiëntnormalisatie?
2. Impliceert de spectrale-normgeometrie een bias naar oplossingen met een lage rang (nucleaire-norm minimalisatie)?

De auteurs stellen dat beide incompleet zijn. In plaats daarvan stellen zij voor dat Muon een flat-update-spectrum bias induceert, waarbij een gelijke amplitude wordt toegewezen over alle actieve singuliere richtingen. Het artikel streeft ernaar deze bias theoretisch en empirisch te karakteriseren, specifiek in de context van ondergedetermineerde regressie waar de optimizer een eindige interpolant moet selecteren uit een affiene manifold van nul-verlies oplossingen.

Methodologie
De studie maakt gebruik van een combinatie van continu-tijd dynamische systemen analyse, variatierekening en gecontroleerde empirische experimenten.

1. Theoretisch Kader:

   • Projected Polar Flow: De auteurs analyseren de continu-tijd flow $\dot{W} = -\eta P_\perp P(W)$, waarbij $P(W)$ de polaire factor van $W$ is en $P_\perp$ de orthogonale projectie is op de raakruimte van de interpolatie-manifold $\mathcal{M} = \{W : XW = Y\}$.
   • Shared-Frame (SF) Aanname: Een structurele conditie waarbij de gradiënt- en gewichtsmatrices dezelfde singular frames delen ($U_G = U_W, V_G = V_W$) wordt geïntroduceerd om de analyse te vereenvoudigen. Stellingen worden bewezen zonder deze aanname voor de self-polar flow, terwijl de SF-aanname de kloof overbrugt naar de werkelijke gradiënt-gestuurde Muon-flow.
   • Functionalen: De analyse maakt gebruik van een paarwijze spectrale functionaal $R_{pw}(W) = -\sum_{i \neq j} \log(\sigma_i + \sigma_j)$ en de schaalinvariante versie $\tilde{R}_{pw}$ om de spectrale vorm te volgen.
   • Ruisgeometrie: Een afleiding van de kritieke batchgrootte wordt uitgevoerd met behulp van een één-stap geliniariseerde signaal-ruis crossover, waarbij de Fréchet-afgeleide van de polaire kaart wordt geanalyseerd.

2. Empirische Validatie:

   • Matrix Sensing: Gecontroleerde experimenten op ondergedetermineerde lineaire regressietaken vergelijken Muon met nucleaire-norm minimaliseerders (opgelost via CVXPY) en gradiëntafdaling (gradient descent).
   • Transformer Pretraining: Kleinschalige pretraining van NanoGPT (124M) en een gematchte ViT-controle op CIFAR-10 om spectrale dynamiek (stable rank, spectrale entropie) in deep learning regimes te observeren.
   • Ablaties: Tests op de koppeling van weight decay (gekoppeld versus ontkoppeld) en shared-frame misalignment om robuustheid te verifiëren.

Belangrijkste Bijdragen

1. One-Step Spectral-Bias Stelling:
   Onder expliciete uitlijning-aannames bewijzen de auteurs dat de polaire update de unieke één-stap entropie-maximaliserende keuze is onder begrensde updates die de gradiënt singular richtingen volgen maar de huidige gewichtsspectra negeren. Zij stellen een spectrale afvlakking-criterium vast: het genormaliseerde spectrum vlakt af indien en slechts indien de meetgeometrie bij voorkeur de richtingen behoudt die overeenkomen met de grootste singuliere waarden (gekwantificeerd door een teken-criterium involving $\alpha_i q_i$).

2. Exacte Singuliere Waarde Dynamiek:
   Het artikel leidt de exacte dynamiek af voor de geprojecteerde polaire flow, waarbij wordt aangetoond dat singuliere waarden afnemen met snelheden $\dot{\sigma}_i = -\eta \alpha_i$, waarbij $\alpha_i$ afhangt van de uitlijning tussen de linker singuliere vectoren van het gewicht en de nulruimte van de meetmatrix. Dit leidt tot een variationele karakterisering waarbij, bij een vaste Frobenius-norm, de unieke minimizer van de spectrale functionaal een plat spectrum is (alle niet-nul singuliere waarden zijn gelijk).

3. Afbakening van Bestaande Verklaringen:

   • Normalisatie: Experimenten tonen aan dat norm-gematchte gradiëntafdaling het gedrag van Muon niet reproduceert.
   • Nucleaire Norm: Theoretische en empirische resultaten laten zien dat de limietpunten van Muon niet overeenkomen met nucleaire-norm minimaliseerders. Terwijl nucleaire-norm minimalisatie een spectrale concentratie (lage rang) bevoordeelt, bevoordeelt Muon spectrale afvlakking (hoge stable rank).

4. Regime-Afhankelijke Bruikbaarheid:
   Het artikel stelt vast dat de flat-spectrum bias van Muon niet universeel superieur is. In kleine NanoGPT pretraining behoudt Muon de stable rank en verbetert het de validatieverlies ten opzichte van AdamW. Echter, in een gematchte kleine ViT-controle keert de rangorde om, wat suggereert dat de bias alleen gunstig is wanneer veel spectrale richtingen actief moeten blijven.

Resultaten

• Theoretisch: Stelling 1 bewijst dat de geprojecteerde self-polar flow het systeem naar een plat spectrum drijft als de meetgeometrie aan een specif kind teken-criterium voldoet. Stelling 2 begrenst de afwijking wanneer de Shared-Frame aanname slechts benaderd wordt. Stelling 4 biedt een formule voor de kritieke batchgrootte $B_{crit}$ in de polar-map SGD, die de gevoeligheid van de polar-map voor ruis koppelt aan $S(\mu) = \sum_{i \neq j} (\sigma_i + \sigma_j)^{-2}$.
• Empirisch (Matrix Sensing): Muon convergeert naar oplossingen met significant hogere nucleaire normen (1.29x–2.02x) en hogere spectrale entropie vergeleken met de minimale nucleaire-norm interpolant, wat de hypothese dat Muon de nucleaire norm minimaliseert weerlegt.
• Empirisch (Transformers): In NanoGPT behoudt Muon een hogere per-layer spectrale entropie en stable rank vergeleken met AdamW, dat de neiging heeft de rang te laten instorten in taal-taken. Daarentegen presteert de ViT-controle beter met AdamW, wat de "regime-afhankelijke" conclusie ondersteunt.
• Ruisgevoeligheid: De afgeleide formule voor de kritieke batchgrootte laat zien dat $B_{crit}$ wordt bepaald door de gevoeligheid van de polar-map voor ruis, die wordt opgeblazen door kleine of bijna-degeneratieve singuliere waarden. Platte spectra bij een vaste Frobenius-straal minimaliseren deze gevoeligheid.

Betekenis en Claims
Het artikel claimt de eerste structurele karakterisering te bieden van de impliciete bias van Muon als een flat-update-spectrum mechanisme in plaats van een low-rank of eenvoudige normalisatie-effect.

• Bescheiden Omvang: De auteurs stellen expliciet dat de resultaten regime-afhankelijk zijn. Muon is geen universele optimizer; de bruikbaarheid komt voort uit spectrale regimes waarin het behouden van activiteit over vele richtingen voordelig is.
• Theoretische Limieten: De stellingen zijn lokaal en structureel, gebaseerd op aannames zoals eenvoudige singuliere waarde-gaps en specifieke meetgeometrieën. Het artikel claimt niet het succes van Muon in alle grootschalige LLM-trainingsscenario's te verklaren, waarbij wordt opgemerkt dat de reductie van transformer-lagen naar het bestudeerde ondergedetermineerde regressiemodel een "conditionele observatie" blijft die onbewezen aannames over de activatie-rangen vereist.
• Praktisch Inzicht: Het werk biedt een diagnostisch instrument voor praktijkgebruikers: als een toepassing vereist om de stable rank en brede spectrale activiteit te behouden, is de flat-spectrum bias van Muon voordelig; als het probleem inherent low-rank of goed geconditioneerd is, kan de bias onnodig of zelfs nadelig zijn.

Het artikel concludeert door open problemen te identificeren, waaronder het karakteriseren van de specifieke meetgeometrieën die het afvlakking-criterium globaal garanderen en het uitbreiden van de ruisgevoeligheidsanalyse naar rechthoekige en rank-deficiënte gevallen.