Technische Samenvatting: MONA – Muon-optimizer met Nesterov-versnelling

1. Probleemstelling

Het trainen van grote taalmodellen (LLM's) vereist in toenemende mate optimalisatoren met superieure steekproefefficiëntie. Hoewel AdamW al meer dan een decennium de standaard is, zijn er recente alternatieven zoals Muon opgekomen. Muon verbetert methoden van de eerste orde door gewichtsmatrices als enkele geometrische eenheden te behandelen en matrixorthogonalisatie (via Newton-Schulz-iteratie) toe te passen op de momentumbuffer. Deze aanpak komt overeen met steilste daling onder de spectrale norm en biedt ongeveer $2\times$ rekenkundige efficiëntie ten opzichte van AdamW bij training van Mixture-of-Experts (MoE) in grote schaal.

Echter, Muon deelt een fundamentele beperking met alle gradiëntmethoden van de eerste orde: het mist expliciete mechanismen om het verlieslandschap te verkennen buiten lokale gradiëntinformatie. Bijgevolg kan Muon vastlopen in scherpe lokale minima, met name in scenario's met grote batches. Scherpe minima worden geassocieerd met slechte generalisatie, terwijl vlakkere minima doorgaans betere downstream-prestaties opleveren. Hoewel de ALTO-optimizer een versnellingsterm introduceerde die gebaseerd is op het exponentiële voortschrijdend gemiddelde (EMA) van gradiëntverschillen om scherpe minima te ontvluchten, vormt zijn ontwerp als een algemene adapter voor Adam en LAMB integratie-uitdagingen met Muon's sterk niet-lineaire orthogonalisatieproces.

2. Methodologie: MONA

De auteurs stellen MONA (Muon Optimizer with Nesterov Acceleration) voor, wat naadloos een krommingsbewuste versnellingsterm integreert direct in Muon's orthogonalisatiekader.

Kernalgoritme

MONA wijzigt de standaard Muon-pijplijn door een versnellingsterm voor momentumaccumulatie en orthogonalisatie in te voeren. Het updateproces voor een matrixparameter $W$ op iteratie $k$ verloopt als volgt:

1. Berekening Gradiëntverschil: Bereken het verschil tussen de huidige stochastische gradiënt $G_k$ en de vorige gradiënt $G_{k-1}$:
   $$D_k = G_k - G_{k-1}$$
2. Update Versnellingbuffer: Onderhoud een exponentieel voortschrijdend gemiddelde van deze verschillen:
   $$A_k = \beta_a A_{k-1} + (1 - \beta_a) D_k$$
   Hierbij regelt $\beta_a$ het geheugen van de versnelling.
3. Krommingsbewuste Gradiënt: Construeer een uitgebreide gradiënt $\tilde{G}_k$ door de geschaalde versnellingsterm toe te voegen aan de ruwe gradiënt:
   $$\tilde{G}_k = G_k + \alpha A_k$$
   De coëfficiënt $\alpha$ schaalt de versnelling. Theoretische motivatie suggereert dat $G_k - G_{k-1}$ benadert $H_k(\theta_k - \theta_{k-1})$, waarbij $H_k$ de Hessiaan is. Aldus wijst de term weg van scherpe minima (waar gradiëntveranderingen groot zijn) en blijft deze klein in vlakke gebieden.
4. Muon-orthogonalisatie: De uitgebreide gradiënt wordt gevoed in de standaard Muon-momentum- en orthogonalisatiestappen:
   $$M_k = \mu M_{k-1} + \tilde{G}_k$$
   $$O_k = \text{Newton-Schulz}(M_k)$$
   $$W_{k+1} = W_k - \eta (\gamma O_k + \lambda W_k)$$

Ontwerprichtlijnen

Het toepassen van de versnellingsterm voor orthogonalisatie is cruciaal. Als deze na orthogonalisatie wordt toegepast, zou dit de orthogonale structuur verstoren waarop Muon vertrouwen voor spectrale-normregularisatie. Door de ruwe gradiënt eerst om te zetten in een krommingsbewuste richting, vangt de momentumbuffer deze verrijkte richtingen op, die vervolgens via Muon's standaardpijplijn worden verwerkt.

Voor vectorwaardige parameters (bijv. embeddings, bias) gebruikt MONA standaard AdamW, in overeenstemming met Muon's bestaande conventies.

3. Theoretische Analyse

Het artikel biedt een convergentieanalyse voor MONA onder standaardaannames (L-gladheid, onbevooroordeelde gradiënten met begrenste variantie, en verwachte directionele uitlijning).

• Convergentiesnelheid: De auteurs bewijzen dat MONA een convergentiesnelheid van $O(K^{-1/2})$ bereikt in niet-convexe settings, wat overeenkomt met de asymptotische snelheden van SGD en AdamW. De versnellingsterm beïnvloedt de constanten in de grens, maar niet de asymptotische snelheid.
• Ontvluchten van Scherpe Minima: De analyse formaliseert hoe de versnellingsterm selectief scherpe minima weerstaat. Nabij een stationair punt is de verwachte versnelling $E[A_k]$ evenredig met de Hessiaan-gewogen som van vorige richtingen. Voor richtingen met grote eigenwaarden (scherpe minima) is de versnellingsgrootte groot, waardoor de optimizer wordt weggeduwd. Voor kleine eigenwaarden (vlakke richtingen) is de versnelling klein, wat stabiele convergentie toelaat. Dit mechanisme behoudt Muon's spectrale-normregularisatie terwijl het krommingsbewustzijn toevoegt.

4. Empirische Resultaten

MONA werd geëvalueerd over drie schalen van MoE-pretrainingsmodellen (1B, 6B en 68B parameters) en een fase van supervised fine-tuning (SFT).

Pretrainingsprestaties

• Modellen: MOE-1B-A0d2B (400B tokens), MOE-6B-A0d5B (1,2T tokens) en MOE-68B-A3B (700B tokens).
• Metingen: Validatieverlies werd gemonitord over domeinen voor code, wiskundig redeneren, algemeen Engels en Chinees academisch tekst.
• Vondsten: MONA bereikte consistent een lager validatieverlies dan zowel Muon als AdamW over alle drie de modelschalen en alle domeinen. Het prestatieverschil was vooral evident in het grotere 68B-model, waar MONA een stabielere convergentie vertoonde.

Downstream Benchmarks (MOE-68B-A3B)

Het 68B-model getraind met MONA werd geëvalueerd op een uitgebreide reeks benchmarks:

• Algemene Capaciteit: MONA behaalde de hoogste gemiddelde score (0,4557) op benchmarks waaronder MMLU, MMLU-Pro, CMMLU, CEVAL en BBH, en presteerde beter dan Muon (0,4478) en AdamW (0,4382).
• Wiskundig Redeneren & Codegeneratie: MONA toonde superieure prestaties op GSM8K, GPQA en diverse codegeneratie-benchmarks (HumanEval+, MBPP+, BigCodeBench, LiveCodeBench). Opmerkelijk is dat MONA aanzienlijke winsten leverde op taken die overdraagbare codevoorstellingen vereisten.
• Supervised Fine-Tuning (SFT): Toen het met MONA voorgetrainde checkpoint SFT onderging op codegegevens, behield het een lager trainingsverlies gedurende het hele proces in vergelijking met de met Muon voorgetrainde baseline. Evaluaties na SFT op de BigCode-suite bevestigden dat met MONA geïnitieerde modellen hogere scores behaalden op 6 van de 7 taken in vergelijking met met Muon geïnitieerde modellen.

Efficiëntie en Implementatie

Om het geheugennadeel van het opslaan van de vorige gradiënt ($G_{k-1}$) en de versnellingbuffer ($A_k$) aan te pakken, introduceerden de auteurs MONA-Lite:

• Quantisatie: Het opslaan van auxiliaire toestanden in bfloat16 (BF16) in plaats van float32.
• Streaming Berekening: Het berekenen van gradiëntverschillen in-place om de noodzaak van een aparte $G_{k-1}$-buffer te elimineren.
• Resultaat: Deze strategieën verminderen de extra geheugenoverhead met ongeveer 75% zonder in te leveren op trainingskwaliteit. Analyse van de rekenkundige overhead toonde aan dat MONA op het niveau van de optimalisatorstap slechts ~1% trager is, zonder praktisch verschil in doorlooptijd van iteratie in vergelijking met Muon.

5. Betekenis en Claims

Het artikel claimt dat MONA succesvol de kloof overbrugt tussen Muon's geometrische efficiëntie en de behoefte aan krommingsbewuste versnelling. De primaire betekenis ligt in:

1. Algoritmische Integratie: Aantonen dat versnellingstermen kunnen worden geïntegreerd in kaders voor matrixorthogonalisatie zonder hun geometrische eigenschappen te verstoren.
2. Schaalbaarheid: Valideren van de methode over een breed scala aan modelgroottes (1B tot 68B parameters) en tokenaantallen (tot 1 biljoen tokens), waarbij consistente verbeteringen worden getoond ten opzichte van zowel AdamW als Muon.
3. Generalisatie: Empirisch bewijs leveren dat het ontvluchten van scherpe minima tijdens pretraining leidt tot betere generalisatie in downstream-taken, met name in complex redeneren en codegeneratie.
4. Praktijktoepasbaarheid: Aanbieden van een geheugenefficiënte variant (MONA-Lite) die de optimizer levensvatbaar maakt voor omgevingen met beperkte middelen.

De auteurs concluderen dat MONA een state-of-the-art optimizer voor training van LLM's in grote schaal vertegenwoordigt, die superieure convergentie en downstream-prestaties bereikt terwijl het de rekenkundige efficiëntie van het Muon-kader behoudt.