Technische Samenvatting: Aurora – Een Leverage-bewuste Spectrale Optimizer

Probleemstelling

Het artikel identificeert een kritiek pathologie in de training van grote taalmodellen met behulp van de Muon-optimizer, een populaire spectrale optimizer die gewichtupdates benadert via de polaire factor van de gradiënt-momentummatrix. Specifiek voor hoge matrices (waarbij het aantal rijen $m$ groter is dan het aantal kolommen $n$, zoals de up- en gate-projectiematrices in MLP-lagen), kunnen Muon-updates zeer niet-uniforme rijnormen vertonen.

Deze niet-uniformiteit creëert een zelfversterkende feedbackloop:

1. Neuronen die overeenkomen met rijen met kleine normen ontvangen persistent kleine updates.
2. Deze neuronen dragen niet betekenisvol bij aan de outputs van het netwerk, waardoor ze effectief "dood" worden.
3. Het artikel levert mechanistisch bewijs dat dit fenomeen gekoppeld is aan leverage scores, waarbij lage rijnormen in de momentumbuffer correleren met lage leverage scores, wat Muon behoudt in plaats van corrigeert.

De auteurs observeren dat bestaande methoden voor rij-normalisatie (bijv. NorMuon) dit probleem weliswaar verzachten door rijen te herschalen, maar dit doen op een manier die de geometrie van de Muon-update verstoort en deze weg beweegt van de ware polaire factor. Deze afweging tussen het afdwingen van uniforme rijnormen en het behouden van de precieze geometrie van de polaire factor is het centrale probleem dat wordt geadresseerd.

Methodologie

1. Analyse van Normalisatie en Polaire Precisie

De auteurs analyseren eerst de spanning tussen polaire factor precisie en rij-normalisatie.

• Polaire Approximatie Fout (PAE): Ze demonstreren dat een hogere precisie bij het berekenen van de polaire factor (via Newton-Schulz iteraties zoals CANS-12) correleert met een lagere validatieverlies.
• Geometrisch Conflict: Ze bewijzen dat het voor hoge matrices wiskundig onmogelijk is om tegelijkertijd kolom-orthogonaliteit (de eigenschap van de polaire factor) en eenheid-rijnormen te voldoen. Post-hoc normalisatie dwingt de update weg van de ware polaire factor, wat significante geometrische fouten introduceert.
• U-NorMuon: Als tussenstap stellen ze U-NorMuon voor, die rij-normalisatie alleen toepast op hoge matrices en staatvolle momentumbuffers voor normalisatie verwijdert. Dit isoleert het voordeel van rij-normalisatie zonder de overhead van verouderde statistieken, en laat zien dat dit beter presteert dan standaard Muon en NorMuon.

2. De Aurora Optimizer

Om het conflict tussen het behouden van de polaire geometrie en het afdwingen van uniforme rijnormen op te lossen, stellen de auteurs Aurora voor.

• Formulering: Aurora formuleert de update als een geconstrueerd optimalisatieprobleem dat de trace van de gradiënt-dot product maximaliseert onder twee restricties:
  1. Links semi-orthogonaliteit: $U^\top U = \eta^2 I$ (behoud van de polaire factor structuur).
  2. Gelijke rijnormen: $\|U_{i,:}\|^2 = \eta^2 \frac{n}{m}$ voor alle rijen $i$.
• Riemannian-Aurora: Ze leiden een theoretische referentie-solver af, Riemannian-Aurora, die deze restricties behandelt als een gezamenlijk manifold (het snijpunt van de Stiefel-manifold en de gelijke-rijnorm-manifold). Hoewel theoretisch precies, is dit computationeel prohibitief vanwege de noodzaak om grote lineaire systemen op te lossen die leverage-restricties bevatten.
• Praktische Aurora: Ze stellen een iteratieve, praktische instantie voor die afwisselend werkt tussen:
  1. Rij-normalisatie: Het herschalen van rijen om normen gelijk te trekken.
  2. Polaire Factor Approximatie: Het toepassen van een Newton-Schulz iteratie om de matrix opnieuw te orthogonaliseren.
     Dit proces wordt gedempt met een parameter $\beta$ om convergentie te garanderen. De auteurs bewijzen dat, onder lokale regulariteitscondities, deze alternerende projectie lineair convergeert naar een oplossing die aan beide restricties voldoet.

Belangrijkste Bijdragen

1. Analyse van Genormaliseerde Spectrale Optimizers: Het artikel demonstreert dat post-hoc eenheid-rij-normalisatie de precisie van de polaire factor verslechtert voor hoge matrices. Het introduceert U-NorMuon, een staatloze variant die normalisatie beperkt tot hoge matrices, en bewijst dat rij-uniformiteit en precieze orthogonalisatie gelijktijdig wenselijk zijn.
2. Studie naar Neuron Death: De auteurs bieden een mechanistische verklaring voor "neuron death" in met Muon getrainde modellen, door dit te koppelen aan lage leverage scores in hoge matrixparameters. Ze verifiëren dit empirisch in hun eigen trainingruns en in open-weight modellen (bijv. Rnj-1, Laguna-XS.2), waarbij zij significante groepen dode neuronen vinden in dense MLP-lagen.
3. Aurora en Riemannian-Aurora: Ze leiden Aurora af als een methode om uniforme rijnormen af te dwingen terwijl strikt aan het Muon theoretische kader wordt voldaan. Ze bieden ook Riemannian-Aurora aan als een high-cost referentie-solver voor het dual-constraint probleem.
4. Uitgebreide Pre-training Verificatie: De auteurs valideren Aurora over 340M en 1.1B parameter schalen, waarbij ze superieure downstream prestaties laten zien vergeleken met Muon, NorMuon en U-NorMuon.

Resultaten

Downstream Prestaties

• Benchmark Scores: Op de 1.1B schaal presteert Aurora significant beter dan de Muon baseline, met een 9.1 punt hogere score op MMLU. Het behaalt ook de hoogste gemiddelde score over een suite van reasoning en kennis benchmarks (ARC, HellaSwag, etc.) op zowel 340M als 1.1B schaal.
• Speedrun: In de modded-nanoGPT speedrun (optimizer track), zette Aurora samen met bestaande methoden (Contra-Muon, flooring) een nieuwe staat van de kunst, waarbij het de doelstelling voor validatieverlies 25 stappen sneller bereikte dan de NorMuon baseline.

Neuron Utilisatie

• Leverage Scores: Modellen getraind met Muon vertonen een bimodale distributie van leverage scores, met een significante modus nabij nul (dode neuronen). In contrast hiermee produceren Aurora en U-NorMuon isotrope leverage score distributies, wat dode neuronen effectief elimineert.
• Open-Weight Verificatie: Analyse van open-weight modellen getraind met Muon (Rnj-1 en Laguna-XS.2) bevestigt de aanwezigheid van dode neuronen in vroege dense lagen (tot wel ~28% in de eerste laag van Laguna-XS.2), wat de generaliseerbaarheid van het probleem valideert.

Efficiëntie en Schaling

• Throughput: In een gedistribueerde setting is de extra computationele overhead van Aurora minimaal, omdat de extra matrixvermenigvuldigingen kunnen worden overlapt met gradiëntcommunicatie. Op een enkele GPU neemt de throughput licht af, maar de wall-clock overhead in gedistribueerde training is bescheiden.
• MLP Breedte Schaling: De prestatiewinst van Aurora over Muon schalt monotoon met de MLP expansiefactor. Naarmate de MLP-lagen breder worden, neemt het voordeel van Aurora toe, wat suggereert dat het bijzonder effectief is voor het trainen van zeer brede netwerken.

Betekenis en Claims

Het artikel claimt dat Aurora een specifieke, empirisch geobserveerde pathologie in de training van taalmodellen oplost — neuron death in hoge matrixparameters — door een optimizer af te leiden die zowel de theoretische geometrie van de Muon-update respecteert als de praktische noodzaak van uniforme rijnormen.

De auteurs benadrukken een bottom-up ontwerpfilosofie: het afleiden van algoritmen op basis van geobserveerde trainingsdynamiek (neuron death, leverage score imbalances), in plaats van louter op theoretische modellen. Ze argumenteren dat architectuur-optimizer co-design, waarbij optimizers op maat worden gemaakt voor specifieke parametertypes (bijv. hoge versus brede matrices), een veelbelovende richting is die standaard "one-size-fits-all" strategieën overtreft.

Ten slotte suggereert het artikel dat Aurora het effectief trainen van zeer brede MLP-lagen mogelijk maakt, een capaciteit die door huidige Muon-gebaseerde training beperkt kan worden door de accumulatie van dode neuronen. Het werk positioneert Aurora niet alleen als een prestatieverbeteraar, maar als een noodzakelijke correctie op de geometrische aannames die ten grondslag liggen aan de huidige spectrale optimalisatie in grootschalige taalmodellen.