Technische Samenvatting: OrScale – Orthogonaliserende Optimalisatie met Schaalvergroting op Basis van Laagsgewijze Vertrouwensratio

1. Probleemstelling

De Muon-optimizer is naar voren gekomen als een praktisch hulpmiddel voor training op grote schaal, waarbij gebruik wordt gemaakt van orthogonaliserende updates om de richting van matrixgewaardeerde stappen te corrigeren via Newton–Schulz-iteraties (een benadering van de polaire factor). Muon lijdt echter aan een kritieke beperking: terwijl het de richting van updates optimaliseert, laat het de grootte onder controle van een enkele, in wezen globale leersnelheid.

Deze globale schaalvergroting is suboptimaal omdat verschillende lagen (bijvoorbeeld attention-projecties, MLP-up/down-projecties, embeddings) zeer verschillende gradiëntverdelingen en effectieve krommingen vertonen. Hoewel de Moonlight-variant een statische factor per vorm introduceerde om de RMS per element te egaliseren, kan deze niet adaptief reageren op de dynamische drift van laagschalen tijdens de training.

Eerdere pogingen om Layer-wise Adaptive Rate Scaling (LARS) en Layer-wise Adaptive Moment Scaling (LAMB) naar Muon te porteren – die succesvol stapgroottes ontkoppelden van globale leersnelheden in SGD en Adam – zijn mislukt. Het artikel identificeert dat naïeve porten instorten vanwege drie specifieke structurele redenen:

1. Gedegenereerde Nominatoren: Het gebruik van de Frobenius-norm van de orthogonaliserende richting ($\|Q_\ell\|_F$) als noemer voor de vertrouwensratio faalt omdat deze bijna constant is, bepaald door de matrixvorm (rang) en niet door de grootte van de update.
2. Eenheidsmismatch & Clip-Verzadiging: Het gebruik van de ruwe momentum-norm ($\|\tilde{M}_\ell\|_F$) introduceert een eenheidsmismatch (gradiënt-eenheden versus gewichtseenheden), waardoor de vertrouwensratio bijna universeel de bovenste clipgrens ($r_{max}$) bereikt, wat de adaptatie ineffectief maakt.
3. Ontkoppelde Weight-Decay Runaway: Het combineren van gekalibreerde nominatoren met ontkoppelde weight decay creëert een positieve feedbacklus waarbij de gewichtsnormen onbeheersbaar groeien (10–50$\times$), wat de training destabiliseert.

2. Methodologie

De auteurs stellen OrScale voor, een raamwerk gebaseerd op één ontwerpprincipe: De noemer van een laagsgewijze vertrouwensratio moet de Frobenius-norm meten van de echte update-richting die de optimizer van de parameters zal aftrekken.

Kernprincipe

Laat $D_{\ell,t}$ de werkelijke richting in de parameterruimte zijn die wordt afgetrokken van de gewichten $W_{\ell,t}$ voordat de globale leersnelheid wordt toegepast. De vertrouwensratio wordt gedefinieerd als:
$$r_{\ell,t} = \frac{\|W_{\ell,t}\|_F}{\|D_{\ell,t}\|_F + \epsilon}$$
Cruciaal is dat $D_{\ell,t}$ zowel de weight-decay-term als de orthogonaliserende richting moet bevatten.

Twee Variant

Het raamwerk levert twee specifieke implementaties op:

1. OrScale (Algemeen/Visie Standaard):

   • Update-richting: $D_{\ell,t} = \lambda W_{\ell,t} + Q_{\ell,t}$ (waarbij $\lambda$ weight decay is en $Q$ de orthogonaliserende richting).
   • Vormfactor: $s_\ell = 1$ (geen statische schaalvergroting).
   • Kalibratie: Geen.
   • Weight Decay: Direct gekoppeld in de update-stap die is geschaald met de vertrouwensratio.

2. OrScale-LM (LLM Standaard):

   • Update-richting: $D_{\ell,t} = \lambda W_{\ell,t} + s_\ell Q_{\ell,t}$.
   • Vormfactor: Neemt de statische factor van Moonlight over: $s_\ell = 0.2\sqrt{\max(m_\ell, n_\ell)}$ om de RMS per element van AdamW te matchen.
   • Kalibratie: Introduceert een lazy per-laag kalibratieconstante ($c_{\text{denom}, \ell}$) die één keer bij initialisatie wordt berekend. Dit verankert de initiële vertrouwensratio exact op 1 ($r_{\ell,0}=1$).
   • Betekenis van Kalibratie: Dit stelt OrScale-LM in staat om de voor AdamW of Moonlight (en bij uitbreiding AdamW) afgestemde hyperparameters voor leersnelheid en weight decay te erven zonder een nieuwe hyperparameter-zweep te vereisen. Het zorgt voor breedte-invariantie bij initialisatie.

Algorithmische Details

• Gekoppelde Weight Decay: Weight decay is geïntegreerd in de berekening van de vertrouwensratio ($\lambda W$ is onderdeel van de noemer) in plaats van toegepast als een aparte terugtrekking. Dit voorkomt de runaway-feedbacklus die bij ontkoppelde varianten wordt waargenomen.
• Efficiëntie: De methode voegt verwaarloosbare overhead toe (<1% wandkloktijd), waarbij slechts één extra scalair per matrixlaag en twee berekeningen van de Frobenius-norm per stap nodig zijn, wat wordt gedomineerd door de bestaande Newton–Schulz-iteraties.

3. Belangrijkste Bijdragen

• Algoritme: Een unificerende, parameter-arme recept (Algoritme 1) dat laagsgewijze schaalvergroting op basis van vertrouwensratio succesvol integreert in Muon. De auteurs identificeren het unieke "werkzame hoekje" in de ontwerpruimte (noemer van echte update-richting + Gekoppelde WD).
• Foutanalyse: Een rigoureuze mapping van de ontwerpruimte voor Muon-vertrouwensratio's (Tabel 1) die analytisch en empirisch drie faalmodi karakteriseert (gedegenereerde noemer, clip-verzadiging en runaway-groei van gewichten), en aantoont waarom deze falen en waarom de voorgestelde oplossing slaagt.
• Theorie:
  • Convergentie: Bewijst een garantie voor niet-convexe convergentie met een $O(1/\sqrt{T})$-snelheid voor de nucleaire norm.
  • Adaptieve Win: Vestigt een laag-adaptieve afstapconstante $\kappa_{\text{eff}} > 1$ onder meetbare laag-heterogeniteit, wat strikt verbetert ten opzichte van Muon's homogene constante.
  • Scheiding: Bewijst een strikte scheiding van MuTrust/MuScale onder omstandigheden van clip-verzadiging.
  • Kalibratie: Demonstreert dat OrScale-LM de $\mu$P-stijl overdracht van leersnelheid bij initialisatie behoudt.
• Empirica: Omvattende evaluatie over taken voor visie- en taalmodellen.

4. Experimentele Resultaten

Visie (CIFAR-10 / DavidNet)

• Prestatie: OrScale behaalde 94,05% $\pm$ 0,08 validatie top-1 nauwkeurigheid, en eindigde als eerste van vijf baselines.
• Verbetering: Dit vertegenwoordigt een verbetering van +0,35 punten ten opzichte van standaard Muon (93,70%) en +0,30 ten opzichte van Muon+Moonlight (93,75%).
• Stabiliteit: OrScale toonde een breder stabiel venster voor leersnelheid in vergelijking met Muon, consistent met theoretische voorspellingen van laag-adaptieve winsten.
• Ablatie: Een ablatiestudie naar faalmodi bevestigde dat varianten die overeenkwamen met de geïdentificeerde faalmodi (bijvoorbeeld het gebruik van $\|Q_\ell\|_F$ of ruw momentum) de voorspelde symptomen vertoonden (bijvoorbeeld 99,5–100% clip-verzadiging of runaway-groei van gewichtsnormen) en niet konden concurreren met de prestaties van OrScale.

Grote Taalmodellen (Pre-training op FineWeb-Edu)

• Schaalvergroting: Geëvalueerd over vier modelgroottes (125M tot 1,1B parameters).
• Vergelijking: OrScale-LM presteerde beter dan Muon+Moonlight op drie van de vier schalen (125M, 545M, 1,1B) en versloeg AdamW op elke schaal.
• Verliesreductie: Op de 1,1B-schaal behaalde OrScale-LM een finale validatie cross-entropy van 2,6251, vergeleken met 2,6360 voor Muon+Moonlight en 2,7304 voor AdamW.
• Efficiëntie: Het voordeel bleef behouden over het bereik van rekenkracht, waarbij OrScale-LM een duidelijke overwinning liet zien in het regime waar rekenkracht het belangrijkst is voor LLM's.

5. Betekenis en Claims

Het artikel beweert dat OrScale succesvol de kloof overbrugt tussen de spectrale geometrie van Muon en de laagsgewijze adaptiviteit van LARS/LAMB. De primaire betekenis ligt in:

1. Oplossen van het Grootteprobleem: Het biedt de eerste tijdsvariabele, laagsgewijze groottecontrole voor Muon die robuust is tegen structurele faalmodi.
2. Praktische Overdraagbaarheid: Via de OrScale-LM-variant en lazy-kalibratie stelt het practitioners in staat om voor AdamW of Moonlight afgestemde leersnelheden direct over te dragen naar OrScale-LM zonder opnieuw af te stemmen, waardoor de drempel voor adoptie verlaagt.
3. Theoretische Grondslagen: Het biedt een rigoureuze theoretische rechtvaardiging voor waarom eerdere hybride pogingen faalden en bewijst dat de voorgestelde methode strikt betere convergentieconstanten oplevert onder heterogeniteit.
4. Empirische Validatie: Het demonstreert dat deze theoretische verbeteringen vertalen naar tastbare winsten in zowel kleine visietaken als pre-training van grote LLM's, en OrScale vestigt als een state-of-the-art optimizer voor matrixgewaardeerde lagen.

De auteurs tonen bescheidenheid wat betreft de LLM-resultaten, en merken op dat het voordeel op de 1,1B-schaal vergelijkbaar is met ruis van een enkele seed op kleinere schalen, en dat de resultaten daarom moeten worden beschouwd als richtinggevend bewijs van de effectiviteit van de methode in plaats van een definitieve statistische superioriteit in alle regimes. Echter, de consistente verbetering over schalen en de succesvolle ablatie van faalmodi ondersteunen het voorgestelde ontwerpprincipe sterk.