Adam and Its Variants

Adam tracks two exponential moving averages (EMAs) of the gradient $g_t = \nabla_\theta L(\theta_{t-1})$ at each step $t$: the first moment $m_t$, an estimate of $\mathbb{E}[g_t]$, and the second raw moment $v_t$, an estimate of $\mathbb{E}[g_t^2]$ (elementwise, per parameter). The first moment plays the role of classical momentum — it smooths out noise and lets consistent gradient directions accumulate speed. The second moment plays the role of AdaGrad/RMSProp scaling — it estimates how large and how noisy the gradient has typically been for each coordinate, so that parameters with large or erratic gradients get a smaller effective step and parameters with small, consistent gradients get a relatively larger one. The subtlety that distinguishes Adam from a naive combination of momentum and RMSProp is bias correction. Because $m_t$ and $v_t$ are initialized at zero and updated with a decay factor $\beta_1$ (for $m$) and $\beta_2$ (for $v$) close to 1, the raw EMAs are systematically biased toward zero, especially in the first few steps when there is little history to average over. Adam divides by $(1-\beta_1^t)$ and $(1-\beta_2^t)$ respectively to produce bias-corrected estimates $\hat m_t$ and $\hat v_t$ that are unbiased in expectation under a stationary-gradient assumption. Skipping this correction (as an early, naive combination of momentum + RMSProp effectively does) causes the first steps of training to take abnormally tiny steps, because $v_t$ starts near zero and the ratio $m_t/\sqrt{v_t}$ behaves erratically. The final update rule divides the bias-corrected momentum by the square root of the bias-corrected second moment, plus a small constant $\epsilon$ for numerical stability:

This single line encodes both ideas at once: the numerator is a momentum-smoothed gradient direction, and the denominator is a per-coordinate adaptive learning rate that shrinks steps in directions of historically large or noisy gradient magnitude.

Adam is a first-order, stochastic optimizer: at every step it only needs the current mini-batch gradient $g_t$, plus the two running-average buffers carried over from the previous step. There is no need to store or invert any curvature matrix — the "adaptive scaling" is purely diagonal and elementwise, which is what keeps it cheap enough for models with billions of parameters.

Step-by-step mechanism: 1. Initialize parameters $\theta_0$, and set $m_0 = 0$, $v_0 = 0$, $t = 0$. 2. At each step, sample a mini-batch and compute the gradient $g_t = \nabla_\theta L_t(\theta_{t-1})$. 3. Update the first moment (momentum): $m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t$. 4. Update the second moment (uncentered variance): $v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t^2$ (elementwise square). 5. Bias-correct both: $\hat m_t = m_t / (1-\beta_1^t)$, $\hat v_t = v_t / (1-\beta_2^t)$. 6. Update parameters: $\theta_t = \theta_{t-1} - \alpha\, \hat m_t / (\sqrt{\hat v_t} + \epsilon)$. 7. Repeat until convergence or a fixed number of steps. Typical defaults, from the original paper and confirmed as robust defaults across countless later architectures, are $\alpha = 0.001$, $\beta_1 = 0.9$, $\beta_2 = 0.999$, $\epsilon = 10^{-8}$. ``` function ADAM(params θ, loss_fn, α=0.001, β1=0.9, β2=0.999, ε=1e-8, weight_decay λ=0): m ← zeros_like(θ) # first moment buffer v ← zeros_like(θ) # second moment buffer t ← 0 while not converged: t ← t + 1 batch ← sample_minibatch() g ← gradient(loss_fn, θ, batch) # ∇θ L(θ) on this batch if λ > 0 and variant == "Adam-L2": g ← g + λ θ # L2 penalty folded into gradient (plain Adam) m ← β1 m + (1 - β1) g # elementwise v ← β2 v + (1 - β2) (g ⊙ g) # elementwise square m_hat ← m / (1 - β1^t) # bias correction v_hat ← v / (1 - β2^t) if variant == "AdamW": θ ← θ - α ( m_hat / (sqrt(v_hat) + ε) + λ θ ) # decoupled decay else: θ ← θ - α m_hat / (sqrt(v_hat) + ε) # plain Adam step return θ ``` Every operation above — the moving averages, the elementwise square root, the division — is applied per parameter, independently. There is no cross-parameter coupling, which is exactly what makes the whole computation $O(d)$ per step for $d$ parameters and trivially parallelizable across a parameter tensor on a GPU.

Time complexity. Each step performs a constant number of elementwise operations (multiply, add, square, square root, divide) over vectors of size $d$, the total parameter count. This is $O(d)$ per step, identical in order to plain SGD; the constant factor is larger (roughly 4–5x the elementwise arithmetic of vanilla SGD) but this is dwarfed in practice by the cost of the forward/backward pass through the network, which is typically $\Omega(d)$ or much more for convolutional and attention layers. Over $T$ steps, total optimizer overhead is $O(Td)$.

Space complexity. Adam must store $m$ and $v$ alongside $\theta$, i.e. two extra copies of the parameter vector, giving $O(d)$ extra memory — a real practical cost (3x the parameter memory instead of 1x for SGD, or 2x for SGD with momentum) that matters at billion-parameter scale, motivating memory-efficient variants such as Adafactor.

Why bias correction is exactly right, not just a heuristic. Assume the true gradient at each step is drawn i.i.d. (or at least stationarily) with $\mathbb{E}[g_t] = g$. Unrolling the EMA recursion,

Taking expectations under stationarity,

So $\mathbb{E}[m_t] \ne g$ unless $t \to \infty$: the raw moving average systematically underestimates the true mean by a factor of $(1-\beta_1^t)$, which is severe when $t$ is small (e.g. at $t=1$, $\beta_1=0.9$, the factor is only $0.1$ — a 10x understatement). Dividing by $(1-\beta_1^t)$ exactly cancels this factor, giving $\mathbb{E}[\hat m_t] = g$: an unbiased estimator under the stationarity assumption. The identical argument applies to $v_t$ and $(1-\beta_2^t)$. This is the precise mathematical justification for the correction — it is not an ad hoc rescaling but the unique factor that removes the geometric-series bias.

Why the adaptive denominator helps. $\sqrt{\hat v_t}$ estimates the root-mean-square magnitude of recent gradients for that coordinate. Dividing by it approximately normalizes the step so that every coordinate takes a step of comparable scale regardless of how large or small its raw gradients are — this is exactly the AdaGrad/RMSProp insight, inherited by Adam. Coordinates with sparse, rare, but informative gradients (e.g. an embedding row hit once every 10,000 steps) do not get drowned out by coordinates with dense, large gradients (e.g. shared bias terms); their smaller $v$ produces a relatively larger effective step size when they do fire.

Convergence. Adam's original convergence proof (in the online convex optimization setting) was later shown to contain a subtle flaw: a fixed $\beta_2$ can, in adversarial sequences, prevent the regret bound from holding, because $v_t$ can decrease even as gradients grow, letting the effective learning rate increase at the wrong time. The fix used by the follow-up variant AMSGrad is to keep a running maximum of $\hat v_t$ so the denominator is monotonically non-decreasing, restoring a clean regret bound. In practice, plain Adam with standard hyperparameters converges reliably on essentially all standard deep learning benchmarks, and the pathological non-convergent sequences are rarely encountered; AMSGrad is used selectively rather than as a universal replacement.

Tradeoff, honestly stated. Adam's adaptivity buys robustness to learning-rate misspecification and fast early progress, especially valuable for the noisy, non-convex, highly heterogeneous-curvature landscape of deep networks. The cost is (a) 2x-3x optimizer memory versus SGD, and (b) empirically worse generalization than well-tuned SGD+momentum on some large-scale image classification tasks, an effect attributed partly to the $L_2$-vs-weight-decay coupling problem that AdamW addresses below.

Trace two steps of Adam on a single scalar parameter, to make the bias-correction arithmetic concrete. Suppose $\theta_0 = 1.0$, $\alpha = 0.1$, $\beta_1=0.9$, $\beta_2=0.999$, $\epsilon=10^{-8}$, and the observed gradients are $g_1 = 2.0$, $g_2 = 1.8$ (consistently positive, as if the loss is $\theta^2$-like locally).

Step $t=1$:

Notice the bias correction recovers exactly $\hat m_1 = g_1 = 2.0$ and $\hat v_1 = g_1^2 = 4.0$: with only one sample, the corrected estimate should simply equal that sample, and it does — a sanity check that the correction formula is exact at $t=1$ (since $\sum \beta^{t-i} = 1$ trivially when $t=1$).

Step $t=2$:

Compare to what plain SGD with $\alpha=0.1$ would have done: $\theta_2^{SGD} = 1.0 - 0.1(2.0) - 0.1(1.8) = 0.62$, a much larger, more volatile step. Adam's adaptive denominator (near $\sqrt{g^2}\approx |g|$ here, since gradients are consistently around $1.8$–$2.0$) keeps the effective step close to $\alpha$ in magnitude regardless of the raw gradient scale — this is the adaptive-scaling behavior working exactly as designed: it normalizes step size to be comparable to $\alpha$, not proportional to the raw gradient magnitude.

AdamW contrast. Suppose we add weight decay $\lambda = 0.01$. Plain "Adam-L2" folds $\lambda\theta$ into $g$ before the moment estimates are computed, so the decay term itself gets divided by $\sqrt{\hat v_t}$ — coordinates with large historical gradients have their weight decay shrunk, which is backwards (large-gradient parameters usually need more, not less, regularization pressure). AdamW instead applies $\theta \leftarrow \theta - \alpha\lambda\theta$ as a separate, undivided term added directly to the parameter update, so every parameter decays toward zero at the same rate $\alpha\lambda$ regardless of its gradient history — matching the original intent of weight decay in Weight Decay and L2 Regularization.

Adam was introduced by Diederik P. Kingma and Jimmy Ba as a synthesis of ideas that had been developed separately over the preceding years: Momentum and Nesterov Accelerated Gradient provided the notion of accumulating a velocity vector, while AdaGrad and RMSProp provided the idea of a per-coordinate adaptive learning rate driven by accumulated squared gradients. Adam's specific contribution was combining both moving averages with an explicit, mathematically motivated bias-correction step, which proved decisive for stable behavior at the start of training and for robustness across a very wide range of problems with minimal hyperparameter tuning.

Adam and AdamW are the default optimizers behind the majority of modern deep learning training runs: transformer language models (including large-scale pretraining pipelines), vision transformers, diffusion models, and most fine-tuning workflows use AdamW as the standard choice, typically paired with a warmup-then-decay learning rate schedule. Recommendation systems and NLP models with huge, sparse embedding tables particularly benefit from Adam's per-coordinate adaptivity, since most embedding rows are touched rarely and need individualized step sizes. Frameworks such as PyTorch, TensorFlow, and JAX ship Adam and AdamW as built-in, heavily optimized primitives (including fused CUDA kernels), reflecting how central this optimizer has become to the practical deep learning stack. Understanding its mechanics is also the natural gateway to Hyperparameter Tuning, since $\alpha$, $\beta_1$, $\beta_2$, and the weight-decay coefficient are among the most consequential knobs in any training recipe.

Ask the tutor: Can you walk me through why Adam's adaptive denominator can sometimes lead to worse generalization than SGD, and how gradient noise interacts with $\sqrt{\hat v_t}$ near a sharp minimum?