Denoise First, Orthogonalize Later: Understanding Momentum in Muon via Spectral Filtering

Technische Samenvatting: "Denoise First, Orthogonalize Later: Understanding Momentum in Muon via Spectral Filtering"

Probleemstelling
De Muon-optimizer heeft onlangs een sterke empirische prestatie getoond bij het trainen van grote taalmodellen (LLM's) door een momentummatrix ongeveer te orthogonaliseren met behulp van een Newton–Schulz-iteratie. Echter, de theoretische rol van momentum binnen Muon blijft onduidelijk. Bestaande analyses of strippen momentum weg om spectrale updates in isolatie te bestuderen, of behouden het zonder de empirische voordelen uit te leggen. Bovendien suggereert de theoretische literatuur over Muon vaak dat zwaar momentum de convergentiesnelheid kan verslechteren of geen voordeel biedt, wat in strijd is met de praktische observaties waar momentum essentieel is. Daarnaast ontbreeft een rigoureuze verklaring voor de specifieke volgorde van operaties—het toepassen van momentum vóór orthogonalisatie (Pre-polar) versus orthogonalisatie vóór momentum (Post-polar)—ondanks het feit dat Pre-polar in de praktijk aanzienlijk beter presteert dan Post-polar.

Methodologie
De auteurs stellen een gestructureerd signaal-plus-perturbatie gradiëntmodel voor om Muon te analyseren. Zij gaan ervan uit dat de gradiënt $G_t$ op iteratie $t$ uiteenvalt in een coherente, laag-rangige signaalcomponent $G^{\text{sig}}_t$ en een perturbatie $\Xi_t$ met een begrensde variantie en een gemiddelde van nul (BVMZOS).

De kern van de theoretische bijdrage is de analyse van de momentumbuffer $M_t = \beta M_{t-1} + (1-\beta)G_t$ als een spectrale filter. De auteurs onderzoeken drie pipeline-varianten voor de gewichtsupdate $\Delta W_t$:

1. Pre-polar (Standaard Muon): $\Delta W_t \propto O(M_t)$, waarbij momentum eerst wordt toegepast, gevolgd door de polaire factor $O(\cdot)$.
2. Post-polar (Orthogonal-SGDM): $\Delta W_t \propto \tilde{M}_t$, waarbij de polaire factor eerst op de ruwe gradiënt wordt toegepast, waarna momentum wordt geaccumuleerd.
3. Polar-only: $\Delta W_t \propto O(G_t)$, een baseline zonder momentum.

De analyse steunt op de effectieve venstergrootte $T = 1/(1-\beta)$ en maakt gebruik van instrumenten uit de random matrixtheorie, waaronder Wedin's $\sin \Theta$-stelling voor singular subspace perturbatie, om te kwantificeren hoe momentum de singularwaarden en subspaces van de gradiëntmatrix beïnvloedt.

Belangrijkste Bijdragen

1. Momentum als Spectrale Filter (Theorem 1 & Corollary 1):
   De auteurs bewijzen dat momentum onder hun gradiëntmodel fungeert als een laagdoorlaatfilter die de perturbatie (ruis) onderdrukt terwijl de coherente signaalcomponent behouden blijft. Specifiek, naarmate de effectieve venstergrootte $T$ toeneemt, nemen de staart-singularwaarden (gedomineerd door ruis) af met een snelheid van $O(T^{-1/4})$, terwijl de leidende singularwaarden (signaal) begrensd blijven boven nul. Dit creëert een vergrote spectrale kloof tussen de signaal- en ruiscomponenten. Bijgevolg stemmen de top-$r$ singular subspaces van de momentumbuffer $M_t$ betrouwbaarder overeen met de ware signaal subspaces dan die van de ruwe gradiënt.

2. Niet-commutativiteit en Pipeline-volgorde (Theorem 2):
   Het artikel stelt vast dat momentum en orthogonalisatie niet commuteren.

   • Pre-polar Dominantie: Wanneer momentum eerst wordt toegepast, dempt dit de gradiëntstroom (denoising), waardoor de daaropvolgende polaire factor kan opereren op een matrix met een gestabiliseerde singular subspace structuur. De auteurs bewijzen dat Pre-polar een hogere verwachte signaaluitlijning bereikt dan zowel de Post-polar als de Polar-only baselines.
   • Falen van Alternatieven: Voor Post-polar zorgt het eerst orthogonaliseren van de ruizige gradiënt ervoor dat de amplitude-kloof tussen signaal- en ruisrichtingen wordt "uitgewist". De polaire factor vervangt alle singularwaarden door enen, waardoor het onmogelijk wordt om in de daaropvolgende momentumstappen onderscheid te maken tussen signaal en ruis. Theoretisch falen de Post-polar en Polar-only pipelines om de signaalrichting effectief te herstellen, vooral in scenario's met een lage Signal-to-Noise Ratio (SNR).

3. Kwantitatieve Separatie in Lage SNR (Theorem 3):
   In een rank-1 spiked Gaussian model met een lage SNR tonen de auteurs aan dat de signaaluitlijning van de Polar-only pipeline verdwijnt naarmate de matrixdimensie $m \to \infty$. In tegen plaats blijft de Pre-polar pipeline het doelgerichte signaal herstellen, mits de momentum-venstergrootte $T$ voldoende groot is.

Resultaten
De theoretische bevindingen worden gevalideerd door uitgebreide experimenten:

• Visualisatie van de Spectrale Kloof: Met behulp van stationaire probes op NanoGPT laten de auteurs zien dat het vergroten van de momentumcoëfficiënt $\beta$ de kloof tussen de leidende en de staart-singularwaarden van de momentumbuffer vergroot, wat consistent is met Theorem 1.
• Subspace Betrouwbaarheid: Empirische metingen van subspace uitlijningsfouten laten zien dat de afwijking tussen de singular subspaces van de momentumbuffer en de ware singular subspaces afneemt naarmate $T$ toeneemt, wat Corollary 1 valideert.
• Pipeline Vergelijking: Over synthetische modellen, CIFAR-10 en LLaMA 350M NanoGPT training heen, presteert de Pre-polar pipeline consistent beter dan Post-polar en Polar-only wat betreft signaaluitlijning en validatieverlies. Opvallend is dat de prestaties van Post-polar vlak blijven ongeacht $\beta$, wat de theoretische voorspelling bevestigt dat het orthogonaliseren vóór het middelen niet in staat is om verloren gegane signaalrichtingen te herstellen.
• End-to-End Training: In LLM pretraining taken (NanoGPT en LLaMA 350M) bereikt de Pre-polar pipeline (standaard Muon) een lager validatieverlies vergeleken met de andere varianten, wat de claim ondersteunt dat momentum een essentiële component is voor signaalherstel, en niet louter een versnellingsmethode.

Betekenis
Dit artikel biedt de eerste theoretische rechtvaardiging voor waarom Muon momentum nodig heeft en waarom de volgorde van operaties (denoising vóór orthogonalisatie) cruciaal is. Het herformuleert momentum in matrix-gebaseerde optimizers niet alleen als een acceleratiemechanisme, maar als een spectrale filter die singular subspaces stabiliseert tegen ruis. Het leidende principe "denoise first, then orthogonalize" biedt een ontwerpregel voor toekomstige matrix-bewuste optimizers die niet-lineaire spectrale stappen gebruiken, waarbij gesuggereerd wordt dat het scheiden van signaal en ruis aan de vooravond moet staan van enige spectrale normalisatie of orthogonalisatie. De auteurs merken op dat hoewel hun theorie de voordelen van momentum in Muon verklaart, het ook de beperkingen van huidige theoretische grenzen benadrukt die er niet in slagen deze voordelen te voorspellen, wat wijst op een kloof tussen de standaard optimalisatietheorie en het praktische gedrag van spectrale optimizers.