Convergence Rate Analysis of the AdamW-Style Shampoo: Unifying One-sided and Two-Sided Preconditioning

1. Probleemstelling

Het artikel behandelt de theoretische convergentieanalyse van de AdamW-stijl Shampoo-optimizer, een state-of-the-art algoritme dat de AlgoPerf-neuraalnetwerktrainingscompetitie won. Hoewel Shampoo (en varianten zoals SOAP en Muon) superieure empirische prestaties heeft aangetoond ten opzichte van diagonale preconditioneringsmethoden (zoals Adam) door het gebruik van full-matrix preconditioners, ontbraken of waren de theoretische convergentiegaranties voor niet-convexe optimalisatie beperkt.

Specifiek kende de bestaande literatuur drie hoofdgaten:

1. Niet-convexiteit: De meeste eerdere analyses waren beperkt tot convexe settings.
2. Preconditioneringsomvang: Vorige analyses voor niet-convexe gevallen (bijv. Xie et al., 2026) waren beperkt tot éénzijdige preconditionering, waarbij de complexere en effectievere tweezijdige preconditionering die in praktische implementaties wordt gebruikt, werd genegeerd.
3. Praktische componenten: Bestaande theorieën lieten vaak kritieke praktische kenmerken zoals momentum en ontkoppelde gewichtsdecay (AdamW-stijl) buiten beschouwing, die essentieel zijn voor de training van moderne grote taalmodellen (LLM's).

Het artikel beoogt een rigoureuze convergentiesnelheid vast te stellen voor het algemene AdamW-stijl Shampoo-algoritme, waarbij zowel éénzijdige als tweezijdige preconditionering worden verenigd binnen een niet-convex raamwerk.

2. Methodologie

De auteurs analyseren Algoritme 1, dat de AdamW-stijl Shampoo-optimizer vertegenwoordigt. De update-regel omvat:

• Exponentiële Voortschrijdende Gemiddelden (EMA): Het bijhouden van eerste momenten ($M_k$) en tweede momenten ($L_k, R_k$) van gradienten.
• Tweezijdige Preconditionering: De update-stap is $X_{k+1} = (1-\lambda\eta)X_k - \eta L_{k,\epsilon}^{-\frac{1}{2p}} M_k R_{k,\epsilon}^{-\frac{1}{2q}}$, waarbij $L_k$ en $R_k$ opgehoopte uitwendige producten van gradienten zijn (die $E[GG^T]$ en $E[G^TG]$ benaderen).
• Gecombineerde Exponenten: De parameters $p$ en $q$ (waarbij $1/p + 1/q = 1$) stellen het algoritme in staat om te schakelen tussen tweezijdige ($p, q < \infty$) en éénzijdige ($p=1, q=\infty$ of omgekeerd) preconditionering.

Belangrijke Technische Innovaties:
Om de complexiteit van tweezijdige matrixpreconditionering in niet-convexe settings te hanteren, ontwikkelden de auteurs nieuwe bewijstechnieken:

1. Schatten-p Hölder-ongelijkheid: In plaats van de standaard vector-gebaseerde Hölder-ongelijkheden die bij Adam-analyses worden gebruikt, maken de auteurs gebruik van de Schatten-p Hölder-ongelijkheid om de kernnorm ($\|\cdot\|_*$) van de gradient te begrenzen. Dit is cruciaal omdat de preconditioner werkt op de singuliere waarden van de matrixparameters.
2. Matrix Cauchy-Schwarz-ongelijkheid: Om de spectrale norm van de update-stap ($L^{-1/2p} M R^{-1/2q}$) te begrenzen, hanteren ze een matrix-specifieke Cauchy-Schwarz-ongelijkheid om de interactie tussen de linker- en rechterpreconditioners te behandelen.
3. Spectrale Bodem-analyse: Ze introduceren een voorwaarde (Voorwaarde 4) met betrekking tot de ondergrens van de preconditioners ($\hat{\epsilon}$). Ze betogen dat in de praktijk de effectieve spectrale bodem van de preconditioners aanzienlijk groter is dan de numerieke regularisatie $\epsilon$ (bijv. $10^{-12}$), waardoor de theoretische grenzen in overeenstemming zijn met empirische waarnemingen.

3. Belangrijkste Bijdragen

• Gecombineerde Convergentiesnelheid: Het artikel vestigt de eerste convergentiesnelheid voor AdamW-stijl Shampoo in niet-convexe optimalisatie die zowel éénzijdige als tweezijdige preconditionering omvat.
• Afleiding van Optimale Snelheid: Ze bewijzen dat de gemiddelde verwachte kernnorm van de gradient convergeert met een snelheid van:
  $$\frac{1}{K} \sum_{k=1}^K \mathbb{E}[\|\nabla f(X_k)\|_*] \leq O\left( \frac{\sqrt{m+n} \cdot C}{K^{1/4}} \right)$$
  waarbij $m, n$ de matrixdimensies zijn, $K$ het aantal iteraties is, en $C$ overeenkomt met de constante in de optimale SGD-snelheid.
• Theoretische Rechtvaardiging voor Praktische Settings: De analyse verklaart waarom het instellen van de numerieke regularisator $\epsilon$ op extreem kleine waarden (bijv. $10^{-12}$) de convergentie niet belemmert. De theorie toont aan dat de effectieve preconditionerbodem wordt bepaald door de schaal van gradientruis, niet door de numerieke $\epsilon$.
• Normrelatie: De auteurs tonen aan dat de afgeleide kernnormsnelheid analoog is aan de optimale Frobenius-normsnelheid van SGD ($O(K^{-1/4})$) onder de ideale voorwaarde waarbij $\|\nabla f(X)\|_* \approx \Theta(\sqrt{\min(m,n)}) \|\nabla f(X)\|_F$.

4. Resultaten

• Theoretische Grenzen:
  • De worst-case grens hangt af van $\sqrt{m+n}$, wat aanzienlijk strakker is dan de $\sqrt{mn}$ (totale dimensie $d$) afhankelijkheid die wordt gevonden in sommige diagonale methoden of full-matrix ondergrenzen.
  • De convergentiesnelheid komt overeen met de bekende ondergrens voor niet-convexe stochastische optimalisatie ($O(K^{-1/4})$), tot een factor van $\sqrt{m+n}$.
• Empirische Validatie (GPT-2 Pretraining):
  • Gradient Norm Ratio: Experimenten op GPT-2 pretraining toonden aan dat de ratio $\|\nabla f(X)\|_* / \|\nabla f(X)\|_F$ consistent dicht bij de theoretische bovengrens $\sqrt{\min(m,n)}$ blijft. Dit valideert de aanname dat de kernnormsnelheid in de praktijk effectief equivalent is aan de Frobenius-normsnelheid.
  • Spectrale Bodem: De minimale eigenwaarden van de preconditioners ($L_k, R_k$) bleken orde van grootte groter te zijn dan de numerieke $\epsilon$ ($10^{-12}$) en vergelijkbaar met de ruis-afhankelijke schaal $\hat{\sigma}^2/(m+n)$. Dit bevestigt dat de "onpraktische" aanname van een grotere spectrale bodem standhoudt bij echte LLM-training.

5. Betekenis

• Brug tussen Theorie en Praktijk: Dit werk biedt de eerste rigoureuze theoretische onderbouwing voor de specifieke Shampoo-variant die de AlgoPerf-competitie won, en verklaart waarom het in de praktijk zo goed werkt.
• Vooruitgang ten opzichte van Diagonale Methoden: Het valideert theoretisch het superieure karakter van niet-diagonale (full-matrix) preconditionering ten opzichte van diagonale methoden (zoals Adam) voor matrixgestructureerde parameters, en toont aan dat de door Shampoo vastgelegde kruis-parametercorrelaties leiden tot betere convergentierichtingen.
• Richting voor Hyperparameters: De analyse biedt theoretische rechtvaardiging voor de keuze van hyperparameters (bijv. de relatie tussen momentum $\theta$ en tweede-moment decay $\beta$, en de rol van $\epsilon$), en leidt praktici bij het afstemmen van deze optimizers voor grootschalige modellen.
• Fundament voor Toekomstige Optimizers: Door éénzijdige en tweezijdige preconditionering te verenigen en momentum/gewichtsdecay op te nemen, zet dit raamwerk een nieuwe standaard voor de analyse van geavanceerde second-orde-achtige optimizers in deep learning.

Kortom, het artikel verenigt succesvol de theoretische analyse van Shampoo, en bewijst dat het een optimale convergentiesnelheid bereikt voor niet-convexe problemen terwijl het de complexe matrixstructuren vastlegt die inherent zijn aan moderne deep learning-modellen.