Technische Samenvatting: Convergentiebegrens en Kritieke Batchgrootte van de Muon-optimizer

Probleemstelling

Hoewel de Muon-optimizer sterke empirische prestaties heeft getoond als een potentiële opvolger van standaardoptimizers zoals AdamW, blijft de theoretische onderbouwing ervan onderontwikkeld. Bestaande literatuur mist formele convergentiebewijzen voor Muon, met name in praktische settings met Nesterov-momentum en weight decay. Dit artikel beoogt deze hiaten te overbruggen door een rigoureuze convergentieanalyse te bieden en een ondergrens af te leiden voor de kritieke batchgrootte van Muon.

Methodologie

De auteurs voeren een theoretische analyse uit van Muon over vier praktische configuraties:

1. Zonder Nesterov-momentum en zonder weight decay.
2. Met Nesterov-momentum en zonder weight decay.
3. Zonder Nesterov-momentum en met weight decay.
4. Met Nesterov-momentum en met weight decay.

Theoretisch Kader:

• Convergentieanalyse: De auteurs leiden bovengrenzen af voor de gemiddelde verwachte nucleaire norm van de gradiënt, $\frac{1}{T}\sum_{t=0}^{T-1} \mathbb{E}[\|\nabla f(W_t)\|_*]$. Ze maken gebruik van standaardveronderstellingen van $L$-gladheid en begrensde gradiëntvariantie. Een cruciale innovatie is de analyse van de zoekrichting van Muon, die wordt verkregen door het orthogonaliseren van een momentummatrix $C_t$. De auteurs houden rekening met de rang van de momentumfoutmatrix ($r_1$) en de zoekrichting ($r_2$), in plaats van de volledige parameterdimensie.
• Newton-Schulz Benadering: Omdat Muon in de praktijk de orthogonalisatie benadert via de Newton-Schulz iteratie in plaats van exacte SVD, analyseren de auteurs de impact van deze benaderingsfout ($\delta(k)$) op de convergentie, waarbij zij aantonen dat dit slechts een verwaarloosbare term introduceert.
• Afleiding Kritieke Batchgrootte (CBS): De auteurs modelleren het aantal trainingsstappen $T$ dat nodig is om een doelprecisie $\epsilon$ te bereiken als een functie van de batchgrootte $b$. Door de SFO-complexiteit ($T \times b$) te minimaliseren, leiden zij een ondergrens af voor de kritieke batchgrootte $b^*$.

Experimentele Validatie:
De theoretische voorspellingen worden gevalideerd via een tweeledige strategie:

1. Gecontroleerde Full-Muon Experimenten: Een synthetische Teacher-Student taak waarbij Muon alle parameters bijwerkt, wat strikt overeenkomt met de theoretische aannames.
2. Praktische Schaal Hybride Experimenten: Het trainen van ResNet-18/VGG-16 op CIFAR-datasets en Llama3.1 (320M) op de C4-corpus. In deze settings wordt Muon toegepast op matrixvormige parameters, terwijl AdamW de rest afhandelt, wat de standaard implementatie reflecteert.

Belangrijkste Bijdragen

1. Convergentiebewijzen voor Vier Varianten

Het artikel stelt convergentiegaranties vast voor alle vier de Muon-varianten. De afgeleide bovengrenzen voor de gemiddelde verwachte gradiëntnorm zijn van de orde:
$$O\left(\frac{1}{\eta T} + \sqrt{\frac{(1-\beta)r_1}{b}} + \bar{r}\eta\right)$$
(met lichte variaties voor de Nesterov- en weight decay-termen).

• Significantie: In tegen tegenover gradiëntgebaseerde optimizers waarbij de variantie schaalt met het volledige aantal parameters, schaalt de Muon-grens met de rang $r_1$ van de momentumfoutmatrix. Empirisch bewijs suggereert dat $r_1 \ll n$ in overgeparameteriseerde regimes, wat impliceert dat dit een nauwere theoretische grens oplevert.
• Nucleaire Norm: De analyse gebruikt de nucleaire norm, wat een striktere criterium is dan de Frobenius-norm die in veel vectorgebaseerde analyses wordt gebruikt, wat suggereert dat Muon de intrinsieke matrixstructuur van neurale netwerkparameters respecteert.

2. Rol van Weight Decay en Stabiliteit

De auteurs bewijzen dat het opnemen van weight decay zorgt voor de bijna-zekere begrensdheid van zowel de parameternormen ($\|W_t\|_F$) als de gradiëntnormen ($\|\nabla f(W_t)\|_F$) zonder te vertrouwen op de restrictieve aanname van begrensde gradiënten.

• Stabiliteitsconditie: Zij leiden een noodzakelijke conditie af voor stabiliteit: de leersnelheid $\eta$ en de weight decay coëfficiënt $\lambda$ moeten voldoen aan $\eta \leq 1/\lambda$. Experimentele resultaten bevestigen dat training instabiel wordt wanneer $\eta$ deze drempel overschrijdt.
• Impliciete Regularisatie: Hoewel weight decay de asymptotische convergentiesnelheid niet verbetert qua orde, biedt het impliciete regularisatie door normen te begrenzen, wat de superieure empirische prestaties verklaart.

3. Ondergrens Kritieke Batchgrootte (CBS)

Het paper leidt een ondergrens af voor de kritieke batchgrootte $b^*$, waarbij de afhankelijkheid van hyperparameters wordt getoond:
$$b^*_{\text{Muon}} > \frac{9(1-\beta)(1+\sqrt{2\beta})^2 r_1 \sigma^2}{(1-\lambda)^2 \epsilon^2}$$

• Hyperparameter Afhankelijkheid: De formule laat zien dat:
  • Momentum ($\beta$): Naarmate $\beta \to 1$, verdwijnt de term $(1-\beta)$, wat suggereert dat de kritieke batchgrootte afneemt.
  • Weight Decay ($\lambda$): De factor $1/(1-\lambda)^2$ geeft aan dat het verhogen van weight decay de kritieke batchgrootte vergroot.
  • Rang ($r_1$): De grens schaalt met de rang van de momentumfout, wat de CBS koppelt aan de low-rank structuur van gradiënten.
• Kwalitatieve Schaling: De auteurs benadrukken dat de formule een kwalitatieve schalingswet biedt in plaats van een absolute voorspelling, aangezien deze afhangt van probleem-specifieke constanten ($\sigma^2, r_1, \epsilon$). Echter, de ratio van CBS tussen verschillende Muon-varianten hangt enkel af van $\beta$ en $\lambda$.

Resultaten

Theoretische Bevindingen:

• De convergentiesnelheid komt overeen met standaard niet-convexe SGD-achtige methoden ($O(T^{-1/4})$ onder specifieke hyperparameter-schema's), maar met een nauwere variantieterm door de rang-afhankelijkheid.
• De Newton-Schulz benaderingsfout heeft een beperkt effect op de convergentie en neemt exponentieel af met het aantal iteraties.

Empirische Bevindingen:

• Stabiliteit: Experimenten op ResNet-18/CIFAR-10 valideren de stabiliteitsconditie $\eta \leq 1/\lambda$. Gradiëntnormen en de trainingsverlies worden geminimaliseerd nabij deze grens.
• Convergentiesnelheid: Muon met Nesterov-momentum en weight decay bereikt consistent de snelste afname van verlies en gradiëntnorm vergeleken met AdamW, Momentum SGD en andere Muon-varianten.
• Validatie Kritieke Batchgrootte:
  • Hyperparameter Effecten: Experimenten bevestigen dat het toevoegen van Nesterov-momentum de CBS naar grotere waarden verschuift, en het verhogen van weight decay de CBS verder vergroot. Bijvoorbeeld, het verhogen van $\lambda$ van 0 naar 0,5 resulteerde in een verschuiving van de CBS van $\approx 1024$ naar $\geq 4096$, wat de voorspelde $1/(1-\lambda)^2$ schaling (ratio van 4,0) bevestigt.
  • Rang Schaling: Het variëren van de modelbreedte toonde aan dat de effectieve rang $r_1$ schaalt met de breedte, en dat het product $r_1 \cdot \sigma^2$ relatief constant blijft, wat verklaart waarom empirische CBS-waarden vlak bleven over verschillende breedtes ondanks de theoretische rang-afhankelijkheid.
  • Efficiëntie: Muon vereist over het algemeen minder stappen en een lagere SFO-complexiteit dan AdamW, met name bij kleine tot gemiddelde batchgroottes.

Betekenis en Claims

Het paper claimt de eerste rigoureuze convergentieanalyse van Muon te bieden die tegelijkertijd Nesterov-momentum en weight decay incorporeert, waarmee een gat in het theoretische begrip van deze opkomende optimizer wordt gedicht.

• Theoretisch Inzicht: Het demonstreert dat weight decay niet louter een regularisatietool is, maar een mechanisme dat begrensdheid en stabiliteit waarborgt, waardoor de aanname van begrensde gradiënten wordt geëlimineerd.
• Praktische Begeleiding: De afgeleide CBS-ondergrens biedt bruikbare richtlijnen voor beoefenaars. Het verheldert hoe het afstemmen van $\beta$ en $\lambda$ kwalitatief de optimale batchgrootte beïnvloedt, waarbij gesuggereerd wordt dat configuraties met hoog momentum efficiënt getraind kunnen worden met kleinere batches, terwijl een hogere weight decay grotere batches vereist om de efficiëntie te behouden.
• Validatie: Het werk overbrugt de kloof tussen theorie en praktijk door deze voorspellingen te valideren over computer vision en taalmodellering workloads, waarbij wordt bevestigd dat de theoretische schalingswetten standhouden, zelfs in hybride optimalisatie-settings waar Muon slechts op matrixvormige parameters wordt toegepast.

De auteurs merken bescheiden op dat hun analyse van een enkele matrix de laag-specifieke heterogeniteit niet expliciet modelleert, maar zij argumenteren dat de hoge correlatie van gradiëntruis-schalen over lagen de generaliseerbaarheid van hun bevindingen ondersteunt. Zij erkennen ook dat hoewel Muon in hun experimenten Shampoo overtreft, het afleiden van analoge grenzen voor Shampoo een openstaand probleem blijft.