Technische Samenvatting: LionMuon

Probleemstelling

Het trainen van Large Language Models (LLM's) is een massaal optimalisatieprobleem waarbij de rekentijd per stap (FLOPs) en de convergentiesnelheid de primaire bepalende factoren zijn voor de uiteindelijke rekenkosten. Huidige optimalisatoren staan voor een afweging tussen de kwaliteit van de updates en de computationele efficiëntie:

• Op teken gebaseerde methoden (bijv. Lion, Signum): Deze produceren goedkope, coördinaatgewijze tekenupdates ($W_{t+1} = W_t - \eta \cdot \text{sign}(\hat{G}_t)$). Ze zijn computationeel goedkoop en geheugenefficiënt, maar bieden vaak zwakkere update-richtingen in vergelijking met spectrale methoden.
• Spectrale methoden (bijv. Muon): Deze berekenen een matrix-tekenupdate via Newton-Schulz-iteraties ($W_{t+1} = W_t - \eta \cdot \text{msign}(\hat{G}_t)$). Hoewel ze sterkere spectrale richtingen bieden die vaak leiden tot snellere convergentie, brengen ze aanzienlijk hogere kosten per stap met zich mee door meerdere matrixvermenigvuldigingen en verhoogde communicatie-overhead in gedistribueerde omgevingen (waarbij een volledige matrix all-gather vereist is).

De centrale vraag die wordt aangepakt, is of het mogelijk is de lage kosten van op teken gebaseerde updates te combineren met de hoogwaardige richtingen van spectrale updates om een superieure afweging tussen kosten en kwaliteit te bereiken.

Methodologie: LionMuon

De auteurs stellen LionMuon voor, een optimalisator die op een vaste periode $P$ afwisselt tussen Lion-stijl tekenupdates en Muon-stijl spectrale updates.

Kernalgoritme

1. Afwisselend schema: Het algoritme voert elke $P$-de iteratie een Muon-stap uit (Newton-Schulz spectrale update) en Lion-stijl tekenupdates voor de overige $P-1$ iteraties.
2. Gedeelde momentum: Beide update-types delen een enkele dual-Exponential Moving Average (EMA) momentum-buffer ($M_t$). De richting $\hat{G}_t$ wordt berekend via Lion's interpolatie: $\hat{G}_t = \beta_1 M_{t-1} + (1-\beta_1)G_t$.
3. Efficiëntie van toestand: LionMuon onderhoudt slechts één momentum-buffer, wat overeenkomt met het geheugengebruik van Lion/Signum en exact de helft is van dat van AdamW.
4. Speciale gevallen:
   • $P=1$: Puur Muon.
   • $P=\infty$: Puur Lion (of Signum als $\beta_1=\beta_2$).
   • $P \in (1, \infty)$: De voorgestelde afwisselende methode.
   • SignMuon: Een variant die single-EMA momentum gebruikt ($\beta_1 = \beta_2$) met hetzelfde afwisselende schema.

Theoretische Analyse

De auteurs leveren convergentiegrenzen voor LionMuon onder zwaarstaartige ruis, een voorwaarde die vaak voorkomt bij LLM-training.

• Geïnterpoleerde constanten: De convergentiegrens hangt af van periode-gemiddelde gladheid ($\bar{L}$) en ruisniveaus ($\bar{\rho}$) die interpoleren tussen de constanten van puur Muon (spectrale norm) en puur Lion (oneindigheidsnorm).
• Optimale periode: De theorie suggereert dat een optimale periode $P^*$ bestaat die een afwegingsfactor minimaliseert die verband houdt met de verhoudingen van gladheidsconstanten ($L_\infty / L_2$) en ruisniveaus ($\rho_1 / \rho_{nuc}$).
• Complexiteit: De afgeleide complexiteitsgrenzen tonen aan dat LionMuon onder specifieke configuraties dezelfde nauwkeurigheid kan bereiken als puur Muon of Lion met minder bewerkingen, met name wanneer de afwegingsfactor $\phi(P) < 1$.

Belangrijkste Bijdragen

1. SignMuon: Een afwisselende optimalisator die elke $P$ iteraties een Muon-stap uitvoert en anders een Signum-stijl tekenstap. De auteurs tonen aan dat zelfs deze single-EMA variant de afweging tussen verlies en FLOPs verbetert ten opzichte van puur Muon bij $P=2$.
2. LionMuon: De belangrijkste bijdrage, waarbij het momentum van Signum wordt vervangen door de dual-EMA-regel van Lion. Het behoudt de kostenbesparingen van SignMuon terwijl het de convergentie verder verbetert. Het vereist slechts één extra geheel getal als hyperparameter ($P$) en is een directe vervanging voor Lion.
3. Theoretische grenzen: Scherpe complexiteitsgrenzen onder zwaarstaartige ruis die expliciet de interpolatie modelleren tussen Muon- en Lion-regimes, en zo de rekenoptimalistische periode $P$ voorspellen.
4. Empirische validatie: Uitgebreide experimenten over meerdere datasets (FineWeb, SlimPajama, WikiText-103) en architecturen (GPT-2, LLaMA) op schalen van 124M, 355M en 720M parameters.

Experimentele Resultaten

• 124M Schaal: LionMuon (specifiek met $P=2$) Pareto-dominant over Muon, Lion, Signum en AdamW op alle geteste datasets en architecturen. Het bereikt een lager validatieverlies met ongeveer 5,5% minder totale trainings-FLOPs in vergelijking met puur Muon.
• Schaling (355M & 720M): Het voordeel in convergentiekwaliteit blijft bestaan op grotere schalen. Bij 355M presteren LionMuon en SignMuon ($P=2$) beter dan puur Muon. Bij 720M slaat SignMuon ($P=2$) nog steeds puur Muon, hoewel het dual-EMA-voordeel van LionMuon ten opzichte van SignMuon in deze specifieke ondergetrainde budget niet werd gereproduceerd.
• Gedistribueerde training: De methode biedt structurele communicatiebesparingen. Omdat Lion-stappen elementgewijs en lokaal voor shards zijn, halveert het instellen van $P=2$ de frequentie van de dure all-gather-bewerkingen die vereist zijn voor de Newton-Schulz-stap van Muon, ongeacht de batchgrootte.

Betekenis en Beweringen

Het artikel beweert dat LionMuon een structureel superieure aanpak biedt voor LLM-training door de frequentie van dure spectrale updates te ontkoppelen van het aantal iteraties.

• Efficiëntie: Het behoudt de effectiviteit van Muon-stappen terwijl het de gemiddelde kosten per iteratie verlaagt, en komt overeen met de geheugenefficiëntie van op teken gebaseerde methoden.
• Robuustheid: De methode is robuust over verschillende modelgroottes en datasets, en definieert consequent de Pareto-grens voor verlies versus FLOPs.
• Theoretische uitlijning: De empirische resultaten sluiten aan bij de theoretische voorspellingen met betrekking tot de optimale periode $P$, wat suggereert dat intermediaire waarden (bijv. $P=2$ of $P=5$) vaak de beste balans bieden tussen de dure spectrale richting en de goedkope tekenrichting.

De auteurs concluderen dat LionMuon een praktische, directe uitbreiding is voor moderne LLM-training die strikt verbetert op de huidige state-of-the-art (Muon) wat betreft zowel convergentiekwaliteit als computationele efficiëntie. Toekomstig werk wordt voorgesteld voor adaptieve $P$-selectie en schaling boven de 720M.