Adam: A Method for Stochastic Optimization

Two running averages give every weight its own learning rate.

Adam keeps, for each parameter, a running average of its gradient and a running average of its squared gradient. It steps along the first divided by the square root of the second, which makes each step about the same size no matter how large or small that parameter's gradients are.

A 2014 paper with a six-line algorithm, a default learning rate of 0.001 that often works unchanged, and a convergence proof that turned out to be wrong.

Adam, or a close variant such as AdamW, is the optimizer most large neural networks are trained with, and the one PyTorch and JAX users reach for before they have a reason to pick anything else. Diederik Kingma and Jimmy Ba posted it to arXiv in December 2014 and presented it at ICLR 2015 (the author order, a footnote says, was decided by a coin flip over a Google Hangout). The main text runs nine pages and the algorithm fits in one box.

This page builds the algorithm from plain stochastic gradient descent in five steps, each fixing a specific problem:

1. average the gradient over time, to cancel noise (the first moment);
2. divide each parameter's step by the recent size of its gradients (the second moment);
3. correct both averages for starting at zero (bias correction);
4. read the resulting ratio as a step of about $\alpha$ per parameter;
5. check the paper's convergence proof, which fails on a three-step counterexample.

What gradient descent has to work with

Training minimizes a loss, a single number measuring how wrong the model is, as a function of its parameters $\theta$ (every weight, stacked into one long vector). The gradient $g$ is the vector of partial derivatives of the loss with respect to each parameter; it points in the direction the loss rises fastest. Gradient descent steps the other way:

The minus sign makes it descend, and $\alpha$ is the learning rate. Picture walking downhill in fog: you can feel the slope under your feet and nothing beyond it.

SGD works under two constraints. First, the gradient is an estimate. The loss is an average over the whole dataset, far too expensive to evaluate every step, so each step computes the gradient on a random minibatch of, say, 128 examples. That estimate is unbiased (its average over all possible minibatches is the true gradient), but any single one is noisy, and the noise shrinks only as one over the square root of the batch size: quadrupling the batch halves it. Averaging gradients across steps is a cheaper way to cut noise. Second, the method uses slopes only. Curvature lives in the Hessian, a matrix with one row and one column per parameter; for a model with a billion parameters it would have $10^{18}$ entries, so Adam, like SGD, is strictly first-order and never forms it.

With one global $\alpha$, two failures show up constantly. In a long narrow valley (steep across, shallow along) a step big enough to make progress along the floor is too big for the walls, and the iterate bounces from side to side. And with noisy gradients, a fixed step follows each noisy sample, so a large share of its motion goes in directions the next sample reverses. The first moment addresses the noise and part of the bouncing; the second moment addresses the mismatch of scales.

The first moment: averaging the gradient

If consecutive gradient estimates disagree because of noise, average them. Adam keeps a running average of the gradient, the first moment $m_t$, and steps along it instead of along the latest sample:

This is an exponential moving average (EMA). Each step keeps a fraction $\beta_1$ of the old average and mixes in $1-\beta_1$ of the new gradient, so a gradient from $k$ steps ago carries weight proportional to $\beta_1^k$. With the default $\beta_1 = 0.9$, a gradient from 10 steps back carries $0.9^{10} \approx 0.35$ of the weight of the newest one, and the average has an effective memory of about $1/(1-\beta_1) = 10$ steps. Noise with random sign cancels inside that window; a direction that persists across it survives. In the narrow valley the side-to-side component flips sign every step and averages down, while the component along the floor has the same sign every step and accumulates.

This is the momentum idea, with one difference in scale. Classical heavy-ball momentum keeps a velocity $u \leftarrow \mu u + g$ and steps by $\alpha u$. Feed it a constant gradient of 1 with $\mu = 0.9$ and the velocity settles where $u = 0.9u + 1$, at 10, so the step is ten times the gradient. Adam's weights $\beta_1$ and $1-\beta_1$ sum to one, so the same constant gradient drives $m_t$ to $0.9 + 0.1 = 1$, the gradient itself. The average stays on the gradient's own scale, which the division in the next section relies on. Adam also evaluates the gradient at the current parameters $\theta_{t-1}$; the look-ahead version (Nesterov momentum) applied to Adam is a separate method, Nadam.

In Figure 1 the amber dots are noisy gradient samples whose mean flips sign partway through, standing in for a change of direction during training. The teal curve is $m_t$. Drag $\beta_1$ to 0 and it follows every sample; drag it to 0.98 and it is smooth but takes dozens of steps to register the flip.

The second moment: a step size per parameter

Parameters see gradients of very different sizes. In the paper's CNN experiment, weight sharing gives different layers "vastly different gradients", and with SGD practitioners hand-pick a smaller learning rate for the convolution layers. An embedding row for a rare word gets a nonzero gradient only in the few minibatches that contain the word. A single $\alpha$ that is safe for the large-gradient parameters is too small for the rest.

AdaGrad (Duchi, Hazan & Singer, 2011) gave each parameter its own divisor: the square root of the sum of all its squared gradients so far.

A parameter with a history of large gradients gets a large divisor and a small step; a rarely active one keeps a large step, which is what sparse features need. The constant $\delta \ge 0$ sits outside the square root and only guards against division by zero (the Adam paper's own statement of AdaGrad drops it).

The divisor in (3) is a cumulative sum, so it never decreases, and every parameter's effective step $\alpha/\sqrt{\sum g_i^2}$ shrinks for the whole run whether or not the gradients still carry signal. On a convex problem that shrinking is a convergent schedule. On a long non-convex training run it can stall progress.

RMSProp, which Tieleman and Hinton presented in a 2012 Coursera lecture and never published as a paper, replaces the sum with an exponential moving average of the squared gradient. That average is Adam's second moment:

It has the shape of (2), applied to the elementwise square of the gradient. An average forgets: it tracks the size of a parameter's recent gradients instead of their entire history, so the effective step can grow again when gradients get smaller. The decay rate $\beta_2$ sets the length of that memory, about $1/(1-\beta_2)$ steps, or 1000 at the default 0.999.

In Figure 2 the gradient magnitude drops to a third of its starting value partway through, as happens when training settles into a flatter region. Watch the two curves after the drop.

AdaGrad's step, already small from the early large gradients, keeps sinking. Adam's step climbs back up, and the larger $\beta_2$ is, the more slowly it does so: as $\beta_2 \to 1$ the average stops forgetting and behaves like a sum again. The family section makes that limit exact.

The Adam update

Put the two averages together: the first moment gives the direction, the square root of the second moment gives each parameter's divisor, and a small $\epsilon$ keeps the division finite.

The hats mark bias-corrected versions of $m_t$ and $v_t$, defined two sections down; for the moment read them as $m_t$ and $v_t$. Every operation is elementwise, so (5) is a separate one-dimensional update for each parameter, with $\epsilon$ added outside the square root. The name comes from adaptive moment estimation: $m_t$ and $v_t$ estimate the first moment (the mean) and the second raw moment of each parameter's gradient. The whole algorithm, with the paper's defaults:

# Adam: one step on the parameter vector theta (all ops elementwise)
t     = t + 1                        # t starts at 0, so the first step is t = 1
g     = grad(loss, theta)            # gradient at the current theta
m     = b1 * m + (1 - b1) * g        # 1st moment: running mean of g
v     = b2 * v + (1 - b2) * g * g    # 2nd moment: running mean of g^2
m_hat = m / (1 - b1**t)              # bias correction
v_hat = v / (1 - b2**t)
theta = theta - a * m_hat / (sqrt(v_hat) + eps)
# start: m = v = 0 (same shape as theta), t = 0
# defaults: a = 1e-3, b1 = 0.9, b2 = 0.999, eps = 1e-8

Take one parameter through two steps at the defaults. At $t = 1$ the gradient is $g_1 = 0.1$. Both averages start at zero, so

Bias correction divides $m_1$ by $1-0.9 = 0.1$ and $v_1$ by $1-0.999 = 0.001$, giving $\hat m_1 = 0.1$ and $\hat v_1 = 0.01$, so $\sqrt{\hat v_1} = 0.1$ and

The parameter moves by $\alpha$ itself, against the gradient. Plain SGD would have moved it by $\alpha g_1 = 0.0001$, and by $0.1$ if the gradient had been $100$; Adam moves it by 0.001 in both cases. Now suppose the second gradient is noise pointing the other way, $g_2 = -0.1$:

Two gradients of equal size and opposite sign produce a step about a twentieth of the first one, because the numerator averaged them toward zero while the denominator, built from squares, did not shrink. Had $g_2$ been $+0.1$, $\hat m_2$ would be 0.1 and the step a full $\alpha$ again.

Figure 3 puts three optimizers in the same narrow quadratic valley, starting from the same point with one shared step size each. The condition number $\kappa$ is the ratio of the steep curvature to the shallow one: 1 is a round bowl, 40 a long trough. Press Play to run the race, scrub to study any step, and drag $\kappa$ past 29.

On a quadratic with curvature $k$, one SGD step multiplies the distance to the minimum along that axis by $1 - \alpha k$, and the iterate blows up once $|1-\alpha k| > 1$, that is once $\alpha k > 2$. The figure's SGD step and shallow curvature put that threshold at $\kappa \approx 29.4$. Adam's step on each axis is about $\alpha\cdot\mathrm{sign}(g)$ while the gradient keeps its sign, so the curvature of the axis drops out, and one$\alpha$ works across parameters of very different scale. SGD with a learning rate and schedule tuned to a specific problem can match or beat it; Adam tends to be close without that tuning. (The figures on this page are small simulations written for this page, not experiments from the paper.)

Bias correction

Both averages start at zero, since no gradients exist yet. For the first steps each average is a mix of real gradients and that initial zero, so it reads too low. With $\beta_2 = 0.999$ the second moment takes in only a thousandth of each new squared gradient, and after 100 steps it has reached about 10% of the value it is estimating.

Unrolling (4) writes $v_t$ as a weighted sum of every past squared gradient:

Take expectations. If the squared gradient's expected value is the same at every step, it factors out of the sum, and the remaining weights $(1-\beta_2)\beta_2^{\,t-i}$ form a geometric series that sums to exactly $1-\beta_2^t$:

The term $\zeta$ is zero when the second moment is stationary and stays small otherwise, because the exponential weights make old gradients (the ones most likely to differ) count little. The factor $1-\beta_2^t$ is exact: it is the total weight the average has placed on real gradients. The missing $\beta_2^t$ still sits on the initial zero, which contributes nothing, so the average reads low by exactly that factor. Dividing it out gives estimates that are unbiased from the first step:

The first moment gets the same correction with $\beta_1$; the paper notes the derivation is identical. A thermometer switched on in a warm room behaves the same way: if you know what fraction of the way it has warmed up, you can divide the reading by that fraction and get the room temperature before the reading settles. At $t = 1$ the second-moment correction is $1/(1-\beta_2) = 1000$, the factor that turned $v_1 = 10^{-5}$ into $\hat v_1 = 10^{-2}$ in the worked example. As $t$ grows, $\beta_2^t \to 0$ and both corrections fade to 1.

Skipping the correction does not make early steps timid. Both averages are biased low, but the step divides by the square root of $v_t$, and $v_t$ is biased far more than $m_t$ because $\beta_2$ is closer to 1. For a constant gradient the two averages are exactly $(1-\beta_1^t)\,g$ and $(1-\beta_2^t)\,g^2$, so the uncorrected step is

which is 3.16 at $t=1$, peaks at 6.57 at $t = 12$, and takes thousands of steps to come back down to 1. With correction the same constant gradient gives exactly $\alpha$ at every step. The paper's wording is that without correction the "initial steps" are "much larger", and the effect grows as $\beta_2$ approaches 1, exactly the setting sparse gradients call for.

In steady-gradient mode, drag $\beta_2$ to 0.9999 and the uncorrected peak climbs past 20. Set $\beta_1 = 0$ and the uncorrected first step is $1/\sqrt{1-\beta_2}$, 31.6 at the default $\beta_2$. The second mode belongs to the next section.

The ratio and the step bound

With $\epsilon$ set aside, the step is $\Delta_t = \alpha\,\hat m_t/\sqrt{\hat v_t}$. The numerator estimates the mean gradient $\mathbb{E}[g]$. The denominator estimates the square root of the second raw moment $\mathbb{E}[g^2]$, which is the root-mean-square (RMS) size of the gradient, signal and noise together. It is the variance only when the mean is zero: for a gradient with mean $\mu$ and noise of standard deviation $\sigma$, $\mathbb{E}[g^2] = \mu^2 + \sigma^2$.

Two properties follow from dividing a mean by an RMS. The first is scale invariance: multiply every gradient by a constant $c$, and $\hat m_t$ scales by $c$, $\hat v_t$ by $c^2$, its square root by $c$, and the ratio is unchanged. Redo the worked example with $g_1 = 100$ in place of 0.1 and $\hat m_1 = 100$, $\sqrt{\hat v_1} = 100$, and the step is still $\alpha$. The second is a bound. A mean can never exceed an RMS in magnitude,

which is another way of writing $\mathrm{Var}[g] \ge 0$: averaging lets opposite signs cancel, while the RMS counts every sample at full size. So in ordinary training $\hat m_t/\sqrt{\hat v_t}$ sits between about $-1$ and $+1$, and each step moves a parameter by at most about $\alpha$, however steep or shallow its slope. The paper calls this a trust region around the current parameters, the distance beyond which the current gradient estimate is not trusted. Because $\alpha$ bounds the distance traveled per step, you can often choose it from how far the parameters need to move: if good weights lie within about 1 of the initialization, $\alpha = 10^{-3}$ needs on the order of a thousand consistent steps to get there. (The term is Adam's own; TRPO uses the same words for a KL constraint on a policy, a different mechanism.)

The bound is not exactly $\alpha$. Section 2.1 of the paper states two cases:

At the defaults $1-\beta_1 = 0.1$ is larger than $\sqrt{0.001} \approx 0.0316$, so the first case applies and the worst case is $0.1/0.0316 \approx 3.16\,\alpha$. The paper says this case "only happens in the most severe case of sparsity: when a gradient has been zero at all timesteps except at the current timestep." Switch Figure 5 to one gradient after silence to see it: the corrected step is exactly $\alpha$ if the gradient arrives at step 1, dips to about $0.48\,\alpha$ at step 10, and climbs toward $3.16\,\alpha$ only after thousands of silent steps, when $\hat v$ holds a single squared gradient weighted by $1-\beta_2$. The amber line in this mode is the ceiling $(1-\beta_1)/\sqrt{1-\beta_2}$ itself. Raise $\beta_1$ to 0.99 and it falls to $0.32\,\alpha$, so the maximum becomes $\alpha$ (at step 1), the second case of (8).

The same ratio sets how far a parameter moves when its gradient is mostly noise. The paper calls $\hat m_t/\sqrt{\hat v_t}$ the signal-to-noise ratio (SNR), "with a slight abuse of terminology". For a gradient with mean 1 and noise $\sigma$ the ratio fluctuates around $1/\sqrt{1+\sigma^2}$: 1 with no noise, 0.86 at $\sigma = 0.6$, 0.20 at $\sigma = 5$. Close to a minimum the true gradient shrinks while the minibatch noise does not, so the SNR falls and the steps shrink with it. The paper calls this "a form of automatic annealing": a decreasing step size nobody had to schedule.

At large $\sigma$ the teal marker wanders far from the amber circle between resamples, and can land on the wrong side of zero. The average in the numerator covers about 10 steps, so at $\sigma = 5$ its own noise (standard deviation about 1.15) is as large as the signal of 1.

What the paper tested

Section 6 compares Adam with other optimizers on four problems. For every method the learning rate and momentum were searched over a dense grid and the best setting reported, and all methods started from the same initialization. The results are training-cost curves, not tables, so the claims below are the paper's qualitative readings of its figures.

• Logistic regression (convex), L2-regularized, on MNIST's 784-pixel images, minibatch 128, with $\alpha_t = \alpha/\sqrt{t}$ to match the theory. Adam converged about as fast as SGD with Nesterov momentum, and both beat AdaGrad. On IMDB movie reviews as sparse 10,000-word bag-of-words vectors with 50% dropout noise, AdaGrad beat SGD with Nesterov momentum "by a large margin" and Adam matched AdaGrad: the sparse case where per-parameter step sizes matter.
• Multilayer networks on MNIST: two hidden layers of 1000 ReLU units, minibatch 128. Against SFO, a minibatch quasi-Newton method, Adam made faster progress per iteration and per second of wall-clock time; SFO was 5 to 10 times slower per iteration and failed to converge once dropout made the loss stochastic. With dropout, Adam beat AdaGrad, RMSProp, SGD with Nesterov momentum and AdaDelta.
• Convolutional networks on CIFAR-10 (three stages of 5×5 convolutions and 3×3 max pooling, then 1000 ReLU units). Adam and SGD eventually converged considerably faster than AdaGrad. The authors observed that $\hat v_t$ shrank toward zero after a few epochs and was dominated by $\epsilon$, so here the first moment did most of the work. Adam improved on SGD with momentum only marginally, but it adapted the step per layer where SGD needs per-layer rates chosen by hand.
• Bias correction, on a variational autoencoder (500 softplus hidden units, a 50-dimensional latent), sweeping $\beta_1 \in \{0, 0.9\}$, $\beta_2 \in \{0.99, 0.999, 0.9999\}$ and $\log_{10}\alpha$ from $-5$ to $-1$. Without correction, values of $\beta_2$ close to 1 made training unstable, especially in the first epochs. The best results came from small $1-\beta_2$ with correction, and Adam did as well as or better than the uncorrected variant (RMSProp with momentum) at every setting.

The convergence proof and its counterexample

The paper also proves a convergence result, in the online convex optimization framework of Zinkevich (2003). The setup is a repeated game: at each round the optimizer picks parameters $\theta_t$, then a convex loss $f_t$ is revealed, possibly chosen adversarially. Regret compares the losses incurred with those of the best single fixed parameter $\theta^*$ in hindsight:

If $R(T)$ grows more slowly than $T$, the average regret $R(T)/T$ goes to zero, which is the online version of converging to the best answer. The paper claims $R(T) = O(\sqrt{T})$, the best rate known for this general setting. The theorem does not cover Adam as people run it: it assumes a decaying step $\alpha_t = \alpha/\sqrt{t}$ and a decaying momentum $\beta_{1,t} = \beta_1\lambda^{t-1}$ with $\lambda \in (0,1)$, where practical Adam keeps both constant.

In 2018 Reddi, Kale and Kumar showed the proof is wrong even under those assumptions. Proofs of this kind track how the inverse of the effective step size changes from one step to the next,

where $V_t$ is the diagonal matrix of second-moment estimates. The regret bound sums a series of terms weighted by these differences. When every $\Gamma_t$ is positive semi-definite (the effective step never grows on any coordinate), the sum telescopes: consecutive terms cancel and only the endpoints remain, which the assumptions bound. For SGD with a decaying rate and for AdaGrad that holds, because their divisors only grow. Adam's moving average is designed so that $v_t$ can fall when gradients shrink, which makes the effective step rise and $\Gamma_t$ negative on that coordinate. The original proof assumed positive semi-definiteness anyway, and without it the bound does not follow.

Reddi et al. then built a problem where Adam fails. One parameter lives on $[-1, 1]$. The loss is $f_t(x) = Cx$ on every third step and $-x$ on the other two, with $C > 2$. Over one three-step cycle the gradients sum to $C - 2 > 0$, so the total loss falls as $x$ falls and the best fixed point is $x^* = -1$. They set $\beta_1 = 0$ and a deliberately small $\beta_2 = 1/(1+C^2)$ (not the default 0.999), so $v_t$ forgets within a step or two. The rare gradient $C$ inflates the denominator, so its step toward $-1$ is about $\alpha$rather than $C$ times larger; the denominator then decays, and each of the two frequent gradients of $-1$ gets a near-full step toward $+1$. Net drift per cycle is toward $+1$, the worst point in the interval.

The same paper proposed AMSGrad, which divides by the largest second-moment estimate seen so far:

(In Reddi's algorithm the first moment is not bias-corrected.) The running maximum never decreases, so the effective step never grows, $\Gamma_t \ge 0$, and the telescoping argument goes through. In the counterexample the first large gradient sets a high mark that keeps damping the parameter, and AMSGrad converges to $-1$. AMSGrad's own proof was later found to have a gap and was patched by Tran and Phong (2019). None of this changed what practitioners run: the counterexample depends on a small $\beta_2$ and an adversarial loss sequence, and Adam with default settings kept working on real problems.

AdaMax, AdaGrad and other relatives

Section 7.1 of the paper generalizes the denominator. Adam divides by a scaled L2 norm of a parameter's current and past gradients; replace it with an Lp norm and let $p \to \infty$, and the norm becomes a maximum. This gives AdaMax, whose denominator is an exponentially decayed running maximum of the gradient's absolute value:

with $u_0 = 0$. The maximum needs no bias correction, because the newest $|g_t|$ enters at full weight instead of at weight $1-\beta_2$, and there is no $\epsilon$ in the update. Its step obeys the simpler bound $|\Delta_t| \le \alpha$. The suggested defaults are $\alpha = 0.002$ (twice Adam's), $\beta_1 = 0.9$ and $\beta_2 = 0.999$.

AdaGrad is a limit of Adam. Set $\beta_1 = 0$ (no momentum), let $\beta_2 \to 1$, and replace $\alpha$ with the annealed $\alpha_t = \alpha\, t^{-1/2}$. As $\beta_2 \to 1$ the bias-corrected second moment becomes the plain average of all squared gradients, and the annealed step supplies the missing $\sqrt{t}$:

The right-hand side is the step in AdaGrad's update (3) with $\delta = 0$. The correspondence needs all four ingredients. Without bias correction, $\beta_2$ near 1 leaves $v_t$ near zero and the updates grow without bound, the same effect as the uncorrected curve in Figure 5.

The paper also relates Adam to natural gradient descent, which preconditions the gradient with the inverse of the Fisher information matrix, a measure of how much the model's predicted distribution changes when the parameters move. It says $\hat v_t$ approximates the diagonal of that matrix, and that Adam, dividing by $\sqrt{\hat v_t}$, applies the inverse square root of the diagonal (exponent $-1/2$), a milder correction than natural gradient's full inverse (exponent $-1$). The link is an analogy. What Adam accumulates is the squared gradient at the observed labels, an empirical Fisher, which equals the true Fisher only for a log-likelihood loss with labels drawn from the model itself, roughly the situation near a well-fit optimum.

Two further variants are common. L2 regularization added to the loss puts $\lambda\theta$ inside the gradient, where Adam divides it by $\sqrt{\hat v_t}$, so parameters with large gradients get less regularization. AdamW subtracts the weight decay outside the ratio instead. And Section 7.2 of the paper suggests averaging the parameters themselves for evaluation, since the last iterate is noisy:

corrected the same way, by dividing by $1-\beta_2^t$.

Implementation notes

Adam stores two extra vectors the size of the model, $m$ and $v$. In fp32 that is 8 bytes per parameter: 56 GB of optimizer state for a 7-billion-parameter model, more than the fp16 weights themselves. Mixed-precision training usually keeps an fp32 master copy of the weights too, which is why ZeRO counts Adam's training state at 16 bytes per parameter and shards it across GPUs.

The paper gives a reordering that saves two divisions over the whole parameter vector per step, by folding both corrections into a scalar learning rate:

# Same step, reordered as in Sec. 2 of the paper: fold both corrections into a
# per-step learning rate and skip the two divisions over the whole vector.
a_t   = a * sqrt(1 - b2**t) / (1 - b1**t)
theta = theta - a_t * m / (sqrt(v) + eps_hat)
# identical to Algorithm 1 only if eps_hat = eps * sqrt(1 - b2**t)

At $t = 1$ with the defaults, the folded learning rate is $0.001 \times 0.0316/0.1$, or $3.16\times 10^{-4}$. The two forms differ in where $\epsilon$ lands. Multiply (5) through by $\sqrt{1-\beta_2^t}$ and the constant in the denominator becomes $\epsilon\sqrt{1-\beta_2^t}$, which changes with $t$. A library that implements the reordered form with a fixed $\hat\epsilon$ therefore applies an effectively larger $\epsilon$ early in training (by a factor of 31.6 at $t = 1$ with the default $\beta_2$). This matters only when $\sqrt{v}$ is comparable to $\epsilon$, as in the paper's CNN run.