Technische Samenvatting: Optimistische Dual Averaging verenigt moderne optimalisatoren

Probleemstelling

Optimalisatie in deep learning heeft zich ontwikkeld langs twee primaire assen: geometrie en adaptatie (het ontwerpen van updates met impliciete normen of preconditioners die de modelstructuur weerspiegelen, bijv. Adam, Muon) en de compositie van ingrediënten (het combineren van gradientfeedback, momentum en schedules voor stabiliteit, bijv. Lion, schedule-free methoden). Hoewel veel state-of-the-art optimalisatoren (zoals Muon, Lion, AdEMAMix en NAdam) op het eerste gezicht verschillend lijken, ontbreekt er een verenigd theoretisch kader dat hun onderlinge relatie verklaart en principieel advies biedt voor het afstemmen van hyperparameters.

Een specifiek, hardnekkig uitdaging is het afstemmen van weight decay. Bij training over meerdere epochen fungeert weight decay als een regularisator die de normen van iteraties beperkt. Echter, zelfs bij training binnen één epoch, waar overfitting minder een zorg is, blijft weight decay vaak gunstig. Recente schaalregels suggereren het instellen van weight decay evenredig aan $1/d$ (modeldimensie), maar het ontrafelen van de afhankelijkheid van modeldimensie $d$ versus trainingshorizon $n$ is moeilijk. Bovendien variëren optimale weight decay-waarden aanzienlijk tussen verschillende optimalisatoren (bijv. Lion vereist veel sterkere decay dan AdamW), wat kostbare per-optimalisator afstemming noodzakelijk maakt.

Methodologie: SODA

De auteurs introduceren SODA (Stochastic Optimistic Dual Averaging), een generalisatie van Optimistic Dual Averaging (ODA) die de dual verwerking van gradients en de primaal verwerking van iteraties verenigt.

Theoretisch Kader

SODA generaliseert de standaard dual averaging-update:
$$x_{k+1} \in \arg \min_{x \in X} \gamma_k \langle d_k, x \rangle + h_k(x)$$
door een primaal extrapolatiereeks $(y_k)$ en een optimistische dual variabele $(\bar{m}_{k+1})$ in te voeren. Het algoritme handhaaft:

1. Dual Verwerking: Een op momentum gebaseerde gradientgemiddelde $m_{k+1}$ en een optimistische gemiddelde $\bar{m}_{k+1}$ die de huidige gradient $\nabla f(y_k, \xi_k)$ integreert.
2. Primaal Verwerking: Een iteratie-update $x_{k+1}$ en een geëxtrapoleerd punt $y_{k+1}$, gecontroleerd door parameters $\lambda_k$ en $\bar{\lambda}_k$.

De update-regels zijn:
$$\begin{aligned} m_{k+1} &= (1 - \alpha_k)m_k + \alpha_k \nabla f(y_k, \xi_k) \\ \bar{m}_{k+1} &= (1 - \bar{\alpha}_k)m_{k+1} + \bar{\alpha}_k \nabla f(y_k, \xi_k) \quad (\text{optimisme}) \\ z_{k+1} &\in \partial h_k^*(-\gamma_k \bar{m}_{k+1}) \\ x_{k+1} &= (1 - \lambda_k)x_k + \lambda_k z_{k+1} \\ y_{k+1} &= (1 - \bar{\lambda}_k)x_{k+1} + \bar{\lambda}_k z_{k+1} \quad (\text{primaal extrapolatie}) \end{aligned}$$
waarbij $h_k$ een regularisator (spiegelkaart) is die de geometrie definieert (bijv. $\ell_\infty$ voor Lion, spectrale voor Muon).

Vereniging van Bestaande Optimalisatoren

Het artikel toont aan dat door specifieke parameters en geometrieën te selecteren, SODA meerdere moderne optimalisatoren herstelt als speciale gevallen:

• NAdam / Lion / Muon: Deze worden hersteld onder het "gemoderniseerde" regime waarbij $\bar{\lambda}_k = 0$ (geen primaal extrapolatie voor het punt van gradientevaluatie) en $\bar{\alpha}_k > 0$ (optimisme).
  • Lion: Gebruikt $\ell_\infty$-geometrie ($h_k = K^*$) en specifieke momentumparameters.
  • Muon: Gebruikt spectrale geometrie en Nesterov-momentum ($\bar{\alpha}_k = \alpha_k$).
  • NAdam: Gebruikt een gladgemaakte $\ell_\infty$-geometrie met adaptieve schaling.
• Tabel 1 in het artikel koppelt deze methoden expliciet aan hun corresponderende keuzes voor gradientgemiddelde ($\alpha_k$), optimisme ($\bar{\alpha}_k$) en geometrie ($h_k$).

De SODA Wrapper

Een belangrijke praktische bijdrage is een SODA Wrapper (Algoritme 1) die elke basisoptimalisator (bijv. Adam, Lion, Muon) kan omhullen zonder nieuwe hyperparameters in te voeren.

• Configuratie: Het stelt $\bar{\lambda}_k = 0$ (overeenkomend met het "gemoderniseerde" regime), gebruikt uniforme iteratiegemiddelde $\lambda_k = 1/(k+2)$, en schaalt de stapgrootte $\gamma_k = (k+2)\eta_k$.
• Eliminatie van Weight Decay: De wrapper verankert updates aan de initialisatie $z_0$ in plaats van de oorsprong. De resulterende update-regel implementeert effectief een weight decay-schedule van $\mu_k = 1/(k+2)$.
• Interpretatie: De auteurs betogen dat in dit kader weight decay niet slechts een regularisator is, maar een vorm van primaal middelen. Deze theoretisch onderbouwde $1/k$-decay elimineert de noodzaak om een aparte weight decay-hyperparameter af te stemmen.

Belangrijkste Bijdragen

1. Vereniging via Optimisme en Dualisatie: SODA biedt een enkel sjabloon dat Muon, Lion, NAdam en anderen herstelt als optimistische instanties van Dual Averaging, waardoor hun structurele relaties worden verduidelijkt.
2. Theoretische Waarborgen: Het kader leidt hyperparameterkeuzes af (specifiek de $1/k$-decay-schedule) die theoretisch onderbouwd zijn in convergentieanalyse voor $L$-gladde convexe functies. Het reframed weight decay als een mechanisme voor primaal middelen.
3. Parameter-vrije Wrapper: De voorgestelde SODA wrapper verwijdert de noodzaak om weight decay af te stemmen voor elke basisoptimalisator, terwijl de prestaties behouden of verbeterd blijven.
4. Empirische Validatie: Uitgebreide experimenten tonen consistente verbeteringen over modelgroottes en trainingshorizonten.

Experimentele Resultaten

Experimenten werden uitgevoerd op nanoGPT met de FineWeb100-dataset, waarbij SODA-omhulde optimalisatoren werden vergeleken met afgestemde baselines.

• Overdracht over Horizon: De optimale weight decay voor Muon bleek af te nemen naarmate de trainingshorizon toenam, wat in tegenspraak is met statische schaalregels (zoals $1/d$). De dynamische $1/(k+2)$-decay van de SODA wrapper presteerde consistent beter dan Muon, zelfs wanneer Muon zijn best-afgestemde weight decay kreeg (Figuur 1).
• Algemene Toepasbaarheid: De wrapper verbeterde de prestaties voor Adam, Muon en Scion over verschillende trainingshorizonten (1$\times$ en 4$\times$ Chinchilla-stappen) zonder enige extra afstemming (Figuur 2).
• Schaalbaarheid: Bij het schalen van modelgrootte van 64M naar 1B parameters (volgend op Chinchilla-schaalregels) presteerde SODA (specifiek de optimistische variant SODA$^\dagger$) consistent beter dan de basisoptimalisatoren (Scion/uScion), waarbij de prestatiekloof groter werd naarmate de modelgrootte toenam (Figuur 4).
• Vereenvoudiging: De SODA$^\dagger$-configuratie vereenvoudigde de Muon-baseline door Adam te vervangen door Lion voor alle lagen en het gebruik van één momentumhyperparameter, maar bereikte toch concurrerende of superieure prestaties.

Betekenis en Beweringen

Het artikel beweert dat SODA een principieel mechanisme biedt voor het overdragen van optimalisatiestrategieën over verschillende trainingshorizonten en modelschalen zonder de noodzaak van hyperparameterafstemming.

• Reframing van Weight Decay: De auteurs stellen een verschuiving in perspectief voor: weight decay moet niet alleen worden begrepen als een regularisator tegen overfitting, maar als een vorm van primaal middelen die de optimalisatietrajectorie verankert aan de initialisatie.
• Praktische Impact: Door moderne optimalisatoren te verenigen onder het SODA-kader, bieden de auteurs een "plug-and-play" wrapper die de tijdrovende en dure processen van het afstemmen van weight decay elimineert, wat vooral kritiek is naarmate modelgroottes handmatige afstemming onbetaalbaar duur maken.
• Bescheidenheid: De auteurs erkennen beperkingen, waarbij zij opmerken dat hun theoretische analyse momenteel beperkt is tot het convexe kader (hoewel empirisch informatief voor deep learning) en zich richt op single-epoch training. Zij suggereren dat training over meerdere epochen mogelijk nog steeds een hyperparameter voor regularisatiesterkte vereist vanwege overfittingzorgen.

Kortom, SODA verenigt het ontwerpruimte van moderne optimalisatoren door de lens van optimistische dual averaging en biedt een theoretisch onderbouwde, parameter-vrije methode om trainingsstabiliteit en prestaties over schalen te verbeteren.