Adam: A Method for Stochastic Optimization

Inside the paper

The problem

Optimizing a model with many parameters often means using noisy gradients computed from small samples of data. A practical optimizer needs to make progress without storing or computing a large second-order matrix.

The idea

Adam tracks recent gradients and their squares to set a separate update scale for each parameter. It also corrects a distortion caused by starting those running estimates at zero.

How it works

• At each step, compute a gradient on the current stochastic objective and update an exponentially weighted average of the gradient.
• Update a second exponentially weighted average using the squared gradient, with its own decay rate.
• Correct both averages for their initial bias toward zero, which is particularly important early in training.
• Move each parameter using the corrected average gradient divided by the square root of the corrected squared-gradient average, with a small stabilizing term.

What it showed

On the paper’s logistic regression and neural network experiments, Adam generally made useful progress against the compared optimizers. In the variational autoencoder experiment, removing bias correction caused early instability for settings with slowly decaying squared-gradient averages.

Read it for

Start with Algorithm 1 and Sections 2 and 3, then inspect Figure 4 for the bias-correction experiment. The convergence proof in the appendix can wait until you want the mathematical details.

Written in our own words from the paper itself. A draft the editors have not yet reviewed.