Adam: A Method for Stochastic Optimization

Reading guide

Problem and motivation

Stochastic first-order optimization is the backbone of modern machine learning. However, standard Stochastic Gradient Descent (SGD) applies a uniform step size to all coordinates, making it sensitive to gradient scaling, ill-conditioned curvature, and stochastic noise. While AdaGrad adapts coordinate step sizes using historical squared gradients, its learning rate rapidly decays to zero in deep architectures. RMSProp and AdaDelta address this decay with exponential moving averages, but RMSProp lacks theoretical guarantees in non-stationary settings and neither accounts for initialization bias when moving averages start at zero.

The authors aim to combine the benefits of two paradigms: AdaGrad/RMSProp (adapting to sparse and differently scaled gradients) and classical Momentum (accelerating convergence across flat directions and smoothing noise).

Key ideas

• First and Second Moment Estimation: Adam maintains running exponential moving averages of both the gradient (first raw moment $m_t$) and the squared gradient (second uncentered moment $v_t$).
• Bias Correction: Because $m_t$ and $v_t$ are initialized to vectors of zeros, their standard running averages are biased toward zero—especially in early iterations or when decay rates (eta_1, \beta_2) are near 1. Adam analytically corrects this bias by dividing by $(1-{\beta }_1^t)$ and $(1-{\beta }_2^t)$.
• Invariant to Gradient Rescaling: The effective update direction approximately scales with the signal-to-noise ratio ${\hat{m}}_t/\sqrt{{\hat{v}}_t}$, providing natural step-size bounding and invariance to diagonal scaling of the gradients.
• AdaMax: Generalizing the $L_2$ norm tracking of gradients in $v_t$ to the $L_{\infty }$ norm yields AdaMax, a numerically stable variant that dispenses with second-moment bias correction.

Method (with the key equations)

At time step $t$, given objective function $f_t(\theta )$, gradient $g_t={\nabla }_{\theta }{f}_t({\theta }_{t-1})$, step size $\alpha$, exponential decay rates ${\beta }_1,{\beta }_2\in [0,1)$, and small constant $ϵ$:

1. Update biased first moment estimate: $$m_t={\beta }_1{m}_{t-1}+(1-{\beta }_1)g_t$$

2. Update biased second raw moment estimate: $$v_t={\beta }_2{v}_{t-1}+(1-{\beta }_2)g_t^2$$

3. Compute bias-corrected moment estimates: $${\hat{m}}_t=\frac{m_t}{1-{\beta }_1^t}$$ $${\hat{v}}_t=\frac{v_t}{1-{\beta }_2^t}$$

4. Update parameters: $${\theta }_t={\theta }_{t-1}-\alpha \frac{{\hat{m}}_t}{\sqrt{{\hat{v}}_t}+ϵ}$$

Default hyperparameters recommended by the paper are $\alpha =0.001$, ${\beta }_1=0.9$, ${\beta }_2=0.999$, and $ϵ=10^{-8}$.

Results (with the most important numbers)

• Theoretical Convergence: In an online convex optimization framework with bounded gradients and bounded parameter domain $D$, Adam achieves a theoretical regret bound $R(T)={\sum }_{t=1}^T(f_t({\theta }_t)-f_t({\theta }^{\ast }))=O(\sqrt{T})$, yielding an average regret rate of $O(1/\sqrt{T})$.
• Logistic Regression & Multi-layer Perceptrons: On MNIST classification, Adam converged faster than AdaGrad, RMSProp, and SGD with Nesterov momentum across various architectures, displaying minimal sensitivity to hyperparameter perturbations.
• Convolutional Neural Networks: Evaluated on CIFAR-10 with a multi-layer CNN (with dropout and pooling), Adam exhibited faster training loss reduction and reached lower final loss than RMSProp and AdaDelta.

Limitations and open questions

• Convex vs. Non-Convex Theory: The theoretical regret bounds assume convex functions; theoretical convergence behavior on non-convex deep learning objectives was not formally guaranteed.
• Subsequent Theoretical Flaws: Later work (such as Reddi et al., 2018) proved that Adam's convergence proof had an issue where ${\Gamma }_{t+1}-{\Gamma }_t$ is not necessarily positive semi-definite, occasionally leading to empirical divergence (which spurred variants like AMSGrad).
• Generalization Gap: Empirical studies following this paper noted that while Adam trains faster initially, carefully tuned SGD with momentum often finds solutions that generalize slightly better on certain image classification tasks.