AdamW vs Adam

Vanilla stochastic gradient descent (SGD) computes the gradient $\nabla L(\theta_t)$ on a mini-batch and takes a step $\theta_{t+1} = \theta_t - \alpha \nabla L(\theta_t)$. This has fundamental limitations that become severe for deep networks with millions of parameters:

• Oscillation in ravines: Most loss landscapes are highly elongated — the Hessian eigenvalue ratio (condition number) can exceed $10^4$ for deep networks. The gradient points along the steepest descent direction, which in a ravine is nearly perpendicular to the optimal path. This causes violent zigzagging: each step overshoots along the narrow dimension while making little progress along the long dimension. For a 2D quadratic with eigenvalues $\lambda_1 = 100$ and $\lambda_2 = 1$, the gradient-to-optimal-direction angle is $\arctan(100) \approx 89.4°$ — nearly orthogonal.

• No memory of past gradients: Each step uses only the current mini-batch gradient, which is a noisy estimate of the true gradient (mini-batch variance $\propto 1/B$ for batch size $B$). With batch size 32 on a dataset of 50K samples, each gradient estimate sees only 0.064% of the data. The optimizer has no mechanism to average out this noise across iterations.

• Sensitive to learning rate: The optimal learning rate is $\alpha^* \approx 2/(\lambda_{\max} + \lambda_{\min})$ where $\lambda$ are Hessian eigenvalues. For a poorly conditioned loss ($\lambda_{\max}/\lambda_{\min} = 10^4$), this gives $\alpha^* \approx 2 \times 10^{-4}$ — too small for fast convergence along the flat directions. Any larger $\alpha$ causes divergence along the steep direction; any smaller causes extremely slow progress.

• Equal treatment of all parameters: A single global learning rate is applied to all parameters. But different layers, different weight matrices, and different dimensions of the same tensor may have vastly different gradient magnitudes. In a Transformer, embedding gradients might be $100\times$ smaller than attention logit gradients — a single learning rate cannot be optimal for both.

• Stochastic noise amplification: Mini-batch gradients have variance that doesn't decrease with more training iterations (unlike full-batch gradient descent). SGD's trajectory looks like a random walk near the minimum, unable to converge to the exact optimum without learning rate annealing.

Momentum (Polyak, 1964; widely adopted via Sutskever et al., 2013) adds a "velocity" vector $v_t$ that accumulates an exponential moving average of past gradients. This is the foundation of Adam's first moment estimate $m_t$:

$v_t = \beta v_{t-1} + g_t$ $\theta_t = \theta_{t-1} - \alpha v_t$

where $g_t = \nabla L(\theta_t)$ is the current gradient, $\beta \in [0, 1)$ is the momentum coefficient, and $v_0 = 0$. Unrolling the recursion, $v_t$ is a weighted sum of all past gradients: $v_t = \sum_{i=0}^{t} \beta^{t-i} g_i$ — recent gradients are weighted more heavily (exponential decay with half-life $\approx 1/\ln(1/\beta)$ steps).

Why this solves SGD's problems:

1. Oscillation damping: In a ravine, gradients oscillate between positive and negative along the narrow dimension. The momentum average cancels these oscillations: $\beta \cdot (+g) + (-g) \approx 0$ after a few steps. Along the consistent long-axis direction, gradients accumulate: $v \rightarrow g/(1-\beta) = 10g$ at steady state with $\beta = 0.9$. The effective step size along consistent directions is amplified by $1/(1-\beta) = 10\times$.

2. Noise averaging: The exponential moving average acts as a low-pass filter on the stochastic gradient signal. With $\beta = 0.9$, the effective average is over the last $1/(1-\beta) = 10$ gradients, reducing variance by roughly $10\times$ compared to raw SGD. This is equivalent to using a $10\times$ larger batch size — for free.

3. Escape from saddle points and flat regions: In flat regions where $g_t \approx 0$, accumulated velocity $v_t$ from previous non-zero gradients carries the parameters forward — like a ball rolling through a valley floor. Without momentum, the optimizer would stall in flat regions indefinitely.

4. Standard hyperparameter: $\beta = 0.9$ is the near-universal default, meaning velocity is a weighted average of the last ~10 gradient steps. Nesterov momentum ($v_t = \beta v_{t-1} + g(\theta_t - \alpha \beta v_{t-1})$, evaluating the gradient at a "lookahead" position) gives a small additional improvement by incorporating curvature information, and is the default in PyTorch's SGD(nesterov=True).

Limitation that motivates Adam: Momentum gives the same effective learning rate to all parameters — it addresses the oscillation and noise problems but not the per-parameter scale problem. Parameters with consistently large gradients and parameters with small sparse gradients both use $\alpha v_t$. This is where adaptive methods (RMSProp, Adam) enter the picture.