Muon+: Towards Better Muon via One Additional Normalization Step

Dit artikel introduceert Muon+, een verbeterde versie van de Muon-optimizer die door een extra normalisatiestap na orthogonalisatie consequent betere prestaties boekt bij het vooraf trainen van taalmodellen van verschillende groottes.

De kern

Stel je voor dat je een gigantisch, slim computerbrein (een kunstmatige intelligentie) aan het trainen bent. Dit brein moet miljoenen boeken, artikelen en websites lezen om te leren spreken en denken. Dit proces heet "pre-training".

Het probleem is dat dit trainen extreem moeilijk en duur is. Het computerbrein moet constant nieuwe informatie verwerken, maar soms raakt het in de war. Het leert dingen op een manier die niet optimaal is, alsof het probeert te rennen met zware stenen aan zijn voeten.

In de wereld van AI-wetenschap heet de "trainer" die dit brein helpt leren een optimizer. De bekendste trainer heet Adam, maar er is een nieuwe, slimme trainer opgedoken genaamd Muon.

Wat doet Muon? (De Ordehandhaver)

Muon is al heel goed. Het werkt als een strenge maar slimme leraar die zegt: "Hé, jullie (de leerstof) moeten niet in de war raken. Laten we alles netjes en ordelijk maken."

Technisch gezien zorgt Muon ervoor dat de "momentum" (de kracht waarmee het brein leert) orthogonaal wordt. In gewone taal: het zorgt ervoor dat alle verschillende richtingen waarin het brein leert, perfect loodrecht op elkaar staan, zoals de assen op een grafiek. Dit voorkomt dat het brein in één richting blijft hangen en andere belangrijke richtingen negeert.

Wat is MUON+? (De Extra Schoonmaakbeurt)

De auteurs van dit paper zeggen: "Muon is geweldig, maar we kunnen het nog beter maken."

Ze hebben MUON+ bedacht. Dit is Muon, maar dan met één extra stap aan het einde van elke leerbeurt: normalisatie.

Laten we een analogie gebruiken om dit te begrijpen:

De Analogie van de Orkestdirigent
Stel je voor dat je een groot orkest (het AI-model) hebt dat een symfonie speelt.

1. De oude trainer (Adam): Laat iedereen spelen zoals ze willen. Het klinkt vaak rommelig.
2. Muon: Zorgt ervoor dat elke sectie (viool, blaas, percussie) perfect op elkaar afgestemd is en niet in de weg zit van de andere. De muziek klinkt al veel strakker.
3. MUON+ (De nieuwe trainer): Doet precies wat Muon doet, maar voegt daarna een extra stap toe: het regelen van het volume.

Na het afstemmen van de secties, kijkt MUON+ naar elke muzikant en zegt: "Jij klinkt een beetje te hard, jij een beetje te zacht. Laten we even het volume van iedereen op hetzelfde niveau zetten, zodat het geluid perfect in balans is."

Die "volume-regeling" is de normalisatie. Het zorgt ervoor dat de updates (de veranderingen in het brein) niet te groot of te klein worden, maar altijd in een gezond, gebalanceerd bereik blijven.

Waarom is dit belangrijk?

De onderzoekers hebben dit getest op heel veel verschillende modellen, van kleine tot gigantische (zoals GPT en LLaMA). Ze hebben het getest in twee situaties:

1. De standaard situatie: Het brein leert net genoeg om goed te worden.
2. De "overtraining" situatie: Het brein leert veel meer dan nodig is (alsof je een student laat studeren voor een examen dat al lang voorbij is).

De resultaten:

• Altijd beter: In bijna elke test presteerde MUON+ beter dan de originele Muon. Het model leerde sneller en maakte minder fouten.
• Stabiel: Zelfs als je het brein extreem lang laat trainen (met miljoenen extra woorden), blijft MUON+ stabiel. Het "vermoeit" niet en raakt niet in de war.
• Eenvoud: Het geheim zit hem in de eenvoud. Ze hebben geen ingewikkelde nieuwe wiskunde bedacht, maar gewoon die ene extra "volume-regeling" toegevoegd.

Conclusie

Kortom: MUON+ is een kleine, slimme upgrade voor het trainen van super-intelligente computers. Het is alsof je een goede leraar (Muon) een extra hulpmiddel geeft (de volume-regelaar) om ervoor te zorgen dat het hele klasje perfect in balans is. Hierdoor worden de AI-modellen slimmer, sneller en betrouwbaarder, zonder dat het de ontwikkelaars veel extra werk kost.

Het is een bewijs dat soms de kleinste aanpassingen de grootste impact hebben op de toekomst van kunstmatige intelligentie.

Technische samenvatting

Probleemstelling

Het vooraf trainen (pre-training) van grote taalmodellen (LLMs) zoals GPT en LLaMA wordt steeds duurder en computationally intensiever naarmate de modelgrootte en de dataset omvang toenemen. Hoewel optimizers zoals Adam en AdamW de standaard zijn, hebben ze beperkingen in schaalbaarheid en efficiëntie voor zeer grote modellen.

De Muon-optimizer is recentelijk geïntroduceerd als een veelbelovend alternatief. Muon werkt op matrixniveau in plaats van vectorniveau en forceert orthogonalisatie van de momentum-matrix via Newton-Schulz-iteraties. Dit helpt het "rank collapse" van gradiënten tegen te gaan, wat essentieel is voor stabiel leren in diepe netwerken. Echter, ondanks het succes van Muon, blijft er ruimte voor verbetering in stabiliteit en finale modelkwaliteit, vooral bij extreme schaalvergrotingen en hoge token-tot-parameter (T2P) verhoudingen.

Methodologie: MUON+

De auteurs stellen MUON+ voor, een eenvoudige maar effectieve uitbreiding op de Muon-optimizer. De kern van de methode is het toevoegen van één extra normalisatiestap direct na de orthogonalisatie.

1. Bestaande Muon-update:

   • Berekening van momentum: $M_t = \mu M_{t-1} + (1 - \mu)G_t$
   • Orthogonalisatie: $O_t = \text{Ortho}(M_t)$ (via Newton-Schulz iteratie om de SVD te benaderen).
   • Update: $W_t = W_{t-1} - \eta \cdot \sqrt{m/n} \cdot O_t$

2. De MUON+ Innovatie:

   • Na het berekenen van de orthogonale matrix $O_t$, wordt een normalisatie-operator $\text{Norm}^{(d)}(\cdot)$ toegepast voordat de gewichtsupdate plaatsvindt.
   • De update wordt: $W_t = W_{t-1} - \eta \cdot \sqrt{m/n} \cdot \text{Norm}^{(d)}(O_t)$.

   De auteurs onderzoeken verschillende richtingen voor deze normalisatie ($d$):

   • Kolom-gewijs (Column-wise): Normalisatie per kolom (L2-norm van kolommen).
   • Rij-gewijs (Row-wise): Normalisatie per rij.
   • Gecombineerd: Eerst kolom, dan rij (col_row) of andersom (row_col).

   De normalisatie zorgt ervoor dat de update-magnitudes gestructureerd worden, wat de stabiliteit van de optimalisatie verbetert zonder de complexiteit van de optimizer significant te verhogen.

Belangrijkste Bijdragen

1. Eenvoudige Architecturale Aanpassing: Het introduceren van een enkele normalisatiestap na orthogonalisatie, wat leidt tot aanzienlijke prestatiewinsten zonder ingewikkelde tweede-moment aanpassingen of manifold-projecties.
2. Uitgebreide Experimentele Validatie: Het paper presenteert uitgebreide pre-training experimenten over een breed scala aan modelgroottes (van 60M tot 1B parameters) en architecturen (GPT- en LLaMA-achtige modellen).
3. Analyse van Schaalbaarheid: Evaluatie van de optimizer in zowel "compute-optimal" regimes (T2P ratio ~20) als industriële "overtraining" regimes (T2P ratio ~200).
4. Ablatie Studies: Systematisch onderzoek naar de invloed van leergraden, normalisatierichtingen en verschillende methoden voor polaire benadering (SVD).

Resultaten

De experimentele resultaten tonen consistent aan dat MUON+ de basis Muon-optimizer overtreft:

• Pre-training Perplexiteit (PPL):
  • Bij GPT-modellen (130M - 774M parameters) levert MUON+ een verbetering op in validatie PPL van 0.91 tot 2.02 punten ten opzichte van Muon.
  • Bij LLaMA-modellen (60M - 1B parameters) wordt een consistente verbetering gezien, met PPL-reducties variërend van 0.37 tot 0.61 punten.
• Stabiliteit bij Overtraining:
  • In experimenten met een T2P-ratio van ~200 (72 miljard tokens getraind), behoudt MUON+ zijn voordeel. De prestatieverbetering blijft stabiel en degradeert niet in latere trainingsfasen, wat wijst op superieure schaalbaarheid.
• Ablatie van Normalisatierichting:
  • Normalisatie is de drijvende kracht achter de verbetering. Zowel rij- als kolom-normalisatie werken beter dan de baseline.
  • Gecombineerde normalisatie (col_row en row_col) levert over het algemeen de beste resultaten op.
  • Rij-normalisatie presteert vaak iets beter dan kolom-normalisatie alleen.
• Robuustheid:
  • MUON+ werkt consistent goed met verschillende SVD-benaderingsmethoden (You, Jordan, PolarExpress).
  • De optimizer is minder gevoelig voor de keuze van de leergraad dan de standaard Muon, wat het makkelijker maakt om te tunen.

Significantie

Deze studie is significant omdat het aantoont dat de structuur van de update (specifiek de normalisatie na orthogonalisatie) een cruciale rol speelt in de stabiliteit en efficiëntie van het trainen van grote taalmodellen.

• Efficiëntie: Het biedt een "low-hanging fruit" oplossing: een minimale code-aanpassing die leidt tot meetbare prestatiewinsten zonder extra rekenkosten of geheugenoverhead.
• Theoretisch Inzicht: Het paper suggereert dat de prestatieverbetering van recente Muon-varianten (zoals NorMuon) grotendeels te danken is aan de normalisatiestap zelf, en niet noodzakelijk aan complexere mechanismen zoals adaptieve tweede-moment schaling.
• Praktische Toepassing: Gezien de consistentie over verschillende modelgroottes en architecturen, is MUON+ direct toepasbaar in industriële pre-training pipelines, wat kan leiden tot snellere convergentie en betere eindmodellen bij gelijke rekenkosten.

Kortom, MUON+ bewijst dat een simpele normalisatiestap na orthogonalisatie een fundamentele verbetering biedt voor de optimalisatie van grote neurale netwerken.