Technische Samenvatting: Muown - Controle van Rij-Norm voor Muon-Optimalisatie

Probleemstelling

Muon is naar voren gekomen als een concurrerende optimalisator voor het voorvertrouwen van grote taalmodellen (LLM), waarbij het vaak AdamW overtreft door gebruik te maken van de matrixstructuur van Transformer-lineaire lagen. Echter, recente observaties geven aan dat het gedrag van Muon op schaal gevoelig is voor gewichtsverval en vatbaar voor numerieke instabiliteiten, zoals bfloat16-overlopen en exploderende attention-logits.

Een kritiek probleem dat is geïdentificeerd, is de consistente opwaartse drift van de spectrale norm ($\|W_t\|_{S_\infty}$) van gewichtsmatrices tijdens het trainen wanneer Muon wordt gebruikt zonder ontkoppeld gewichtsverval. Deze drift schendt de spectrale voorwaarde die nodig is voor stabiel feature-learnen ($\|W\|_{S_\infty} = \Theta(\sqrt{m/n})$). Hoewel ontkoppeld gewichtsverval dit kan mitigeren door de gewichtsmatrix uniform te laten krimpen, is het een niet-gerichte ingreep die het volledige spectrum schaalt, ongeacht welke specifieke richtingen de groei aandrijven, wat potentieel de convergentie in vroege trainingsfasen kan vertragen. Het artikel streeft ernaar de specifieke empirische drijver van deze spectrale norm-drift te identificeren en een meer gerichte controlemechanisme voor te stellen.

Methodologie: Muown

De auteurs introduceren Muown, een directe vervanging voor Muon die spectrale norm-drift aanpakt via een herparameterisatiestrategie in plaats van aggregate beperkingen.

1. Spectrale Norm Decompositie

Het kerninzicht is een decompositie van de spectrale norm in twee factoren (Propositie 1):
$$\|W_t\|_{S_\infty}^2 = \|g_t\|_\infty^2 \cdot \lambda_{\max}(P_t C_t P_t)$$
Waarbij:

• $g_t = \|W_t\|_{row}$ de vector is van rij-magnituden (Euclidische normen van rijen).
• $D_t$ de rij-genormaliseerde matrix is ($W_t = \text{Diag}(g_t)D_t$).
• $C_t = D_t D_t^\top$ de Gram-matrix is van de eenheidsnorm-rijen.
• $P_t$ een diagonaalmatrix is van genormaliseerde rij-magnituden.
• $\lambda_{\max}(P_t C_t P_t)$ de rij-coherentie-factor voorstelt.

Empirische analyse op 500M-modellen onthult dat de rij-magnitude-factor ($\|g_t\|_\infty$) de primaire drijver is van spectrale norm-drift, die in lock-step beweegt met $\|W_t\|_{S_\infty}$. Daarentegen blijft de rij-coherentie-factor goed gemanaged en begrensd gedurende de hele training.

2. Herparameterisatie en Optimalisatie

Gebaseerd op deze diagnose, herparameteriseert Muown de gewichtsmatrix $W$ als:
$$W(g, R) = \text{Diag}\left(\frac{g}{\|R\|_{row}}\right) R$$
Hierbij vertegenwoordigt $g \in \mathbb{R}^m$ de expliciete vector van rij-magnituden, en $R \in \mathbb{R}^{m \times n}$ de richtingscomponent.

• Richtingscomponent ($R$): Muon wordt onveranderd toegepast op $R$, waarbij deze wordt behandeld onder de spectrale norm-geometrie (steilste spectrale daling).
• Magnitude-component ($g$): De rij-magnituden $g$ worden gepromoveerd tot expliciete optimalisatorvariabelen. Ze worden bijgewerkt onder de $\ell_\infty$-geometrie, wat de natuurlijke duale norm is voor de operatornorm van diagonaalmatrices.
• Bijwerkregel: De bijwerking voor $g$ maakt gebruik van een Adam-optimalisator (werkend als een gladde SignSGD onder $\ell_\infty$-geometrie) om stochastische ruis te hanteren, terwijl de bijwerking voor $R$ de standaard Muon-procedure volgt met Newton-Schulz-orthogonalisatie.

Deze aanpak wordt impliciet geïmplementeerd binnen de optimalisatorstaat, waardoor de forward-pass van de modelarchitectuur ongewijzigd blijft.

Belangrijkste Bijdragen

1. Diagnose van Spectrale Drift: Het artikel decomposeert de spectrale norm om de maximale rij-magnitude ($\|g_t\|_\infty$) te identificeren als de empirische drijver van spectrale norm-groei onder Muon, onderscheidend van de stabiele rij-coherentie-factor.
2. Muown-algoritme: Een nieuwe optimalisator die rij-magnituden behandelt als expliciete trainbare variabelen die worden bijgewerkt onder $\ell_\infty$-geometrie, terwijl Muon's richtingsupdates behouden blijven. Dit vermijdt de indiscriminate krimp van gewichtsverval en de orthogonaliteitspriors van spectrale herparameterisaties.
3. Convergentiegaranties: De auteurs bewijzen dat Muown optimale niet-convexe convergentiesnelheden bereikt in zowel deterministische als stochastische regimes. De analyse wordt uitgevoerd in een gemengde duale norm ($\ell_1$ voor magnituden, Schatten-1 voor richtingen) die is afgestemd op de bijwerkgeometrieën. De theoretische grens suggereert een vermindering van de effectieve stochastische ruiscoëfficiënt in vergelijking met Muon.
4. Empirische Validatie: Uitgebreide experimenten op FineWeb-Edu over modelgroottes van 124M tot 2,7B parameters tonen consistente verbeteringen aan.

Experimentele Resultaten

• Perplexiteitsverbeteringen: Muown bereikt consistent lagere perplexiteit dan Muon, SOAP, AdamW en Lion over alle geteste modelgroottes (124M tot 2,7B) en leersnelheden.
• Robuustheid voor Hyperparameters:
  • Leersnelheid: Muown verbreedt het plateau van bijna-optimale leersnelheden over verschillende modelbreedtes, waardoor de gevoeligheid voor per-schaal leersnelheidstuning wordt verminderd.
  • Gewichtsverval: Muown vermindert de noodzaak voor gewichtsvervaltuning; het presteert goed zonder gewichtsverval ($\lambda=0$), terwijl Muon zorgvuldige tuning vereist (optimale $\lambda=0.1$) om divergentie of suboptimale prestaties te voorkomen.
• Stabiliteit: Muown voorkomt de systematische spectrale norm-drift die bij Muon wordt waargenomen. In regimes met hoge leersnelheden (bijv. $4 \times 10^{-3}$) divergeert Muon vaak, terwijl Muown stabiel en bijna-optimale blijft.
• Gradienteneigenschappen: Analyse toont aan dat Muown een hogere genormaliseerde effectieve rang van minibatch-gradiënten behoudt en een lagere stochastische ruiscoëfficiënt vertoont gedurende de training in vergelijking met Muon.
• Berekeningskosten: De methode introduceert verwaarloosbare overhead per stap. Wanneer de optimalisator over ranks wordt geshard (gedistribueerde setting), worden de stapstijden voor Muown en Muon ononderscheidbaar. De geheugenoverhead is minimaal (lineair in het aantal rijen), met slechts enkele megabytes per laag.

Betekenis en Claims

Het artikel claimt dat de spectrale norm-drift die bij Muon wordt waargenomen, geen irreducibele eigenschap van de optimalisator is, maar een gevolg van ongecontroleerde rij-magnitude-groei. Door deze drijver te isoleren en te promoveren tot een expliciete optimalisatorvariabele, bereikt Muown:

• Verbeterde Convergentie: Betere perplexiteit en stabiliteit zonder de noodzaak van agressief gewichtsverval.
• Theoretische Grondslagen: Optimale convergentiesnelheden in een geometrie-bewust kader dat expliciet rekening houdt met de gemengde aard van de updates.
• Praktische Efficiëntie: Een directe vervanging die de runtime van Muon in gedistribueerde settings evenaart, terwijl het superieure robuustheid biedt voor hyperparameterkeuzes.

De auteurs positioneren Muown als een oplossing die de kloof overbrugt tussen de theoretische stabiliteitsvereisten van feature-learnen en de praktische behoeften van grootschalig LLM-training, en biedt een meer gericht alternatief voor ontkoppeld gewichtsverval.