🤖 machine learning

Optimistic Dual Averaging Unifies Modern Optimizers

This paper introduces SODA, a generalization of Optimistic Dual Averaging that unifies modern optimizers like Muon and Lion under a single framework and offers a practical wrapper to automatically eliminate weight decay tuning while consistently improving performance across various training scales.

Original authors: Thomas Pethick, Wanyun Xie, Roman Machacek, Volkan Cevher

Published 2026-05-13

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Thomas Pethick, Wanyun Xie, Roman Machacek, Volkan Cevher

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a massive, complex robot (a deep learning model) how to walk. The robot learns by taking steps, but sometimes it stumbles, overshoots, or gets confused by noisy instructions. To help it learn faster and more stably, engineers use "optimizers"—mathematical recipes that decide exactly how big each step should be and in which direction.

For the last decade, researchers have been inventing new, fancy recipes (like Adam, Lion, Muon, and NAdam). Each recipe has its own secret sauce, but they all seem to work well for different reasons. The problem is that tuning these recipes is like trying to bake a perfect cake without a recipe: you have to guess the right amount of "weight decay" (a knob that prevents the robot from getting too crazy) for every single new model size and training duration.

This paper introduces SODA (Stochastic Optimistic Dual Averaging). Think of SODA not as a new cake recipe, but as a universal baking wrapper that you can put over any existing recipe to make it work better automatically.

Here is the breakdown of how it works, using simple analogies:

1. The "Optimistic" Look-Ahead

Most optimizers are like a driver who only looks at the road directly in front of them. They see a bump, react, and turn.
SODA is like a driver who looks slightly ahead. It uses a technique called "optimism" to guess where the road is going to be before it gets there.

The Analogy: Imagine you are walking down a hallway. A standard optimizer waits until you bump into a wall to turn. An "optimistic" optimizer sees the wall coming, leans forward slightly, and turns before you hit it. This prevents the robot from wobbling back and forth, allowing it to move smoother and faster.

2. The "Dual Averaging" Memory

Optimizers also need to remember the past. Do they remember every single step they ever took? Or just the last one?
SODA uses a method called "Dual Averaging." Instead of just reacting to the current step, it keeps a running average of all the "pushes" (gradients) it has felt so far.

The Analogy: Think of a hiker trying to find the lowest point in a foggy valley.
- Standard Optimizer: "I feel a slope down to the left, so I'll take a step left."
- SODA: "I felt a slope left 10 steps ago, a slope right 5 steps ago, and a slope left just now. If I average all those feelings, the true path is actually slightly forward-left."
  By averaging the "feelings" (gradients) over time, SODA filters out the noise and finds the true path more reliably.

3. The "Magic Wrapper" (No More Tuning)

The biggest headache for AI engineers is Weight Decay. This is a setting that acts like a "leash," keeping the robot's steps from getting too wild.

The Old Way: You have to guess the perfect leash length. If the robot is small, you need a short leash. If the robot is huge, you need a long one. If you train for a long time, you need to change the leash length again. It's a constant guessing game.
The SODA Way: The authors discovered that if you wrap any optimizer (like Adam or Muon) with SODA, you don't need to guess the leash length anymore.
- SODA automatically adjusts the leash based on a simple rule: 1 divided by the number of steps taken so far.
- The Analogy: Imagine a leash that automatically shortens itself as you walk further. You don't need to hold the leash; the leash knows exactly how tight to be at step 1, step 100, or step 10,000.

4. What Did They Prove?

The paper claims that SODA unifies many of the best modern optimizers (like Muon, Lion, and NAdam) under one single mathematical umbrella. It shows that these different "fancy" methods are actually just special versions of this same "optimistic averaging" idea.

The Results:

Better Performance: When the authors wrapped popular optimizers with SODA, the models learned faster and reached lower error rates.
No Extra Work: They didn't have to tune any new knobs. They just applied the wrapper.
Beating the Best: In their tests, SODA even beat the "best-tuned" versions of the original optimizers. The original optimizers needed a human to carefully adjust the weight decay to get good results; SODA just did it automatically and did it better.

Summary

Think of SODA as a "smart autopilot" you can attach to any existing car engine (optimizer).

It looks ahead (Optimism) to avoid bumps.
It remembers the journey (Averaging) to stay on course.
It automatically adjusts the brakes (Weight Decay) based on how far you've traveled, so you never have to guess how hard to brake.

The paper concludes that this simple, automatic adjustment makes training large AI models more stable, faster, and less dependent on human guesswork.

Technical Summary: Optimistic Dual Averaging Unifies Modern Optimizers

Problem Statement

Deep learning optimization has evolved along two primary axes: geometry and adaptation (designing updates with implicit norms or preconditioners reflecting model structure, e.g., Adam, Muon) and the composition of ingredients (combining gradient feedback, momentum, and schedules for stability, e.g., Lion, schedule-free methods). While many state-of-the-art optimizers (such as Muon, Lion, AdEMAMix, and NAdam) appear distinct, they lack a unified theoretical framework that explains their relationship and provides principled guidance for hyperparameter tuning.

A specific, persistent challenge is the tuning of weight decay. In multi-epoch training, weight decay acts as a regularizer constraining iterate norms. However, even in single-epoch training where overfitting is less of a concern, weight decay often remains beneficial. Recent scaling rules suggest setting weight decay proportional to $1/d$ (model dimension), but disentangling the dependence on model dimension $d$ versus training horizon $n$ is difficult. Furthermore, optimal weight decay values vary significantly across different optimizers (e.g., Lion requires much larger decay than AdamW), necessitating costly per-optimizer tuning.

Methodology: SODA

The authors introduce SODA (Stochastic Optimistic Dual Averaging), a generalization of Optimistic Dual Averaging (ODA) that unifies the dual processing of gradients and the primal processing of iterates.

Theoretical Framework

SODA generalizes the standard dual averaging update:
$x_{k+1} \in \arg \min_{x \in X} \gamma_k \langle d_k, x \rangle + h_k(x)$
by introducing a primal extrapolation sequence $(y_k)$ and an optimistic dual variable $(\bar{m}_{k+1})$ . The algorithm maintains:

Dual Processing: A momentum-based gradient average $m_{k+1}$ and an optimistic average $\bar{m}_{k+1}$ that incorporates the current gradient $\nabla f(y_k, \xi_k)$ .
Primal Processing: An iterate update $x_{k+1}$ and an extrapolated point $y_{k+1}$ , controlled by parameters $\lambda_k$ and $\bar{\lambda}_k$ .

The update rules are:
$\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{optimism}) \\ 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{primal extrapolation}) \end{aligned}$
where $h_k$ is a regularizer (mirror map) defining the geometry (e.g., $\ell_\infty$ for Lion, spectral for Muon).

Unification of Existing Optimizers

The paper demonstrates that by selecting specific parameters and geometries, SODA recovers several modern optimizers as special cases:

NAdam / Lion / Muon: These are recovered under the "Modernized" regime where $\bar{\lambda}_k = 0$ $\overset{ˉ}{λ}_{k} = 0$ (no primal extrapolation for the gradient evaluation point) and $\bar{\alpha}_k > 0$ $\overset{α}{ˉ}_{k} > 0$ (optimism).
- Lion: Uses $\ell_\infty$ geometry ( $h_k = K^*$ ) and specific momentum parameters.
- Muon: Uses spectral geometry and Nesterov momentum ( $\bar{\alpha}_k = \alpha_k$ ).
- NAdam: Uses a smoothed $\ell_\infty$ geometry with adaptive scaling.
Table 1 in the paper explicitly maps these methods to their corresponding choices of gradient averaging ( $\alpha_k$ ), optimism ( $\bar{\alpha}_k$ ), and geometry ( $h_k$ ).

The SODA Wrapper

A key practical contribution is a SODA Wrapper (Algorithm 1) that can wrap any base optimizer (e.g., Adam, Lion, Muon) without introducing new hyperparameters.

Configuration: It sets $\bar{\lambda}_k = 0$ (matching the "modernized" regime), uses uniform iterate averaging $\lambda_k = 1/(k+2)$ , and scales the stepsize $\gamma_k = (k+2)\eta_k$ .
Weight Decay Elimination: The wrapper anchors updates to the initialization $z_0$ rather than the origin. The resulting update rule effectively implements a weight decay schedule of $\mu_k = 1/(k+2)$ .
Interpretation: The authors argue that in this framework, weight decay is not merely a regularizer but a form of primal averaging. This theoretically grounded $1/k$ decay eliminates the need to tune a separate weight decay hyperparameter.

Key Contributions

Unification via Optimism and Dualization: SODA provides a single template that recovers Muon, Lion, NAdam, and others as optimistic instances of Dual Averaging, clarifying their structural relationships.
Theoretical Guarantees: The framework derives hyperparameter choices (specifically the $1/k$ decay schedule) that are theoretically grounded in convergence analysis for $L$ -smooth convex functions. It reframes weight decay as a mechanism for primal averaging.
Parameter-Free Wrapper: The proposed SODA wrapper removes the need to tune weight decay for any base optimizer while maintaining or improving performance.
Empirical Validation: Extensive experiments show consistent improvements across model sizes and training horizons.

Experimental Results

Experiments were conducted on nanoGPT using the FineWeb100 dataset, comparing SODA-wrapped optimizers against tuned baselines.

Transfer Across Horizon: The optimal weight decay for Muon was found to decrease as the training horizon increases, contradicting static scaling rules (like $1/d$ ). The SODA wrapper's dynamic $1/(k+2)$ decay consistently outperformed Muon even when Muon was given its best-tuned weight decay (Figure 1).
General Applicability: The wrapper improved performance for Adam, Muon, and Scion across different training horizons (1 $\times$ and 4 $\times$ Chinchilla steps) without any additional tuning (Figure 2).
Scaling: When scaling model size from 64M to 1B parameters (following Chinchilla scaling), SODA (specifically the optimistic variant SODA $^\dagger$ ) consistently outperformed the base optimizers (Scion/uScion), with the performance gap widening as model size increased (Figure 4).
Simplification: The SODA $^\dagger$ configuration simplified the Muon baseline by replacing Adam with Lion for all layers and using a single momentum hyperparameter, yet achieved competitive or superior performance.

Significance and Claims

The paper claims that SODA offers a principled mechanism for transferring optimization strategies across different training horizons and model scales without the need for hyperparameter tuning.

Reframing Weight Decay: The authors posit a shift in perspective: weight decay should be understood not just as a regularizer against overfitting, but as a form of primal averaging that anchors the optimization trajectory to the initialization.
Practical Impact: By unifying modern optimizers under the SODA framework, the authors provide a "plug-and-play" wrapper that eliminates the tedious and expensive process of tuning weight decay, which is particularly critical as model sizes make manual tuning prohibitively expensive.
Modesty: The authors acknowledge limitations, noting that their theoretical analysis is currently restricted to the convex setting (though empirically informative for deep learning) and focuses on single-epoch training. They suggest that multi-epoch training might still require a hyperparameter for regularization strength due to overfitting concerns.

In summary, SODA unifies the design space of modern optimizers through the lens of optimistic dual averaging and provides a theoretically grounded, parameter-free method to improve training stability and performance across scales.