Adam: A Method for Stochastic Optimization

Why This Matters

In machine learning, finding the optimal parameters of a neural network is akin to navigating a high-dimensional, foggy mountain range. Standard Stochastic Gradient Descent (SGD) takes steps directly proportional to the slope (gradient). However, this naive approach suffers from severe drawbacks:

• Ravines & Saddle Points: If a landscape is steep in one direction but gentle in another, SGD oscillates wildly across the steep walls while making agonizingly slow progress along the valley floor.
• Uniform Learning Rates: SGD applies the exact same step-size scale to all parameters, ignoring the fact that some features might be rare and require larger updates, while others are frequent and need smaller, precise adjustments.
• Noisy Gradients: Mini-batch approximations introduce high variance, leading to erratic trajectories.

Prior solutions like AdaGrad addressed sparse gradients by scaling updates inversely with the sum of all historical squared gradients, but their learning rates decayed to zero too quickly. RMSProp fixed this with an exponential moving average, but lacked a principled way of handling momentum and initialization bias.

The Big Idea

Adam (Adaptive Moment Estimation) combines the best of both worlds: the ability of AdaGrad to handle sparse gradients and the ability of RMSProp to handle non-stationary objectives. It achieves this by maintaining running estimates of both the mean (first moment) and the uncentered variance (second raw moment) of the gradients:

How It Works

At each timestep $t$, Adam computes the gradient $g_t$ of the objective function. It then updates the biased first and second moment estimates:

Because $m_0$ and $v_0$ are initialized as zero vectors, these running averages are heavily biased towards zero, especially during early steps or when decay rates $\beta_1$ and $\beta_2$ are close to 1. To counteract this, Adam applies bias corrections:

Finally, the parameters $\theta_t$ are updated using the scaled step:

Interactive Optimizer Sandbox

Experience how Adam compares directly to SGD and RMSProp on a 2D anisotropic valley (a classic optimization test landscape where the vertical axis is much steeper than the horizontal axis).

Click anywhere on the canvas to set a new starting point, then press Play to watch the trajectories.

The Power of Bias Correction

One of the critical contributions of Adam is the initialization bias correction. Since $m_t$ and $v_t$ are initialized to zero, they are biased towards zero, especially during early steps when $\beta$ is close to 1.

Without the $1 - \beta^t$ correction term, the step sizes in the first few iterations would be drastically smaller (or larger, depending on the scale of gradients) than they should be, leading to slow starts or instability.

Empirical Results

The authors evaluated Adam across several benchmark datasets and architectures, demonstrating consistent, state-of-the-art convergence rates.

Limitations & Open Questions

While Adam is the default optimizer for most deep learning architectures today, it is not without its limitations:

• Generalization Gap: Despite converging faster during training, models trained with Adam sometimes generalize slightly worse to unseen test data compared to carefully tuned SGD with momentum.
• Non-Convergence in Specific Settings: In 2018, Reddi et al. pointed out that Adam can fail to converge to the optimal solution in certain settings where informative gradients are rare but large. This led to the creation of variants like AMSGrad.
• Hyperparameter Sensitivity: Though Adam is marketed as "robust to hyperparameters", extreme values of $\beta_1$ and $\beta_2$ can cause training to diverge completely, especially in deep reinforcement learning or GANs.

Glossary

Citation

@inproceedings{kingma2015adam,
  title={Adam: A Method for Stochastic Optimization},
  author={Kingma, Diederik P and Ba, Jimmy},
  booktitle={International Conference on Learning Representations (ICLR)},
  year={2015}
}