← All tutorials

Adaptive Optimizers: AdaGrad, RMSProp and Adam

Each weight gets its own step size.

On this page

Each method keeps the one above and adds one change.

1 One learning rate cannot fit every direction

Every figure here uses the loss of the series, L = 5x² + 0.2y², started at (−3.5, 5). I call its shape, a long narrow valley, the ravine; it stands for the ill-conditioned loss of a real network, where one direction is far steeper than another. Its numbers are illustrative, computed by script. Plain SGD at α 0.15 takes this first step.

The steep axis caps α (SGD, Momentum and Nesterov, section 1.1), so the flat axis crawls. The methods below give each coordinate its own step size.

2 AdaGrad

AdaGrad divides each coordinate's gradient by the size of all its past gradients.

At step 1, √G = |g|, so every coordinate steps exactly α. After that the sum only grows, so the steps only shrink. GloVe word vectors were trained with AdaGrad at an initial rate of 0.05 (Pennington et al. 2014): a rare word's row gets few gradients, so its steps stay large.

3 RMSProp

RMSProp replaces AdaGrad's forever-sum with an average that forgets. With a decay rate β of 0.99, the share of the old average kept each step, a value from k steps ago gets this weight.

This one average also keeps batch norm's running statistics. RMSProp applies it to g² and calls the average E. At its first step, with β₂ 0.99 and α 0.5:

RMSProp's memory is about 100 steps at β₂ 0.99 instead of forever. The memory matters when gradient sizes change, as when a balancer cuts one task's weight (Gradient Balancing Inside the Training Loop, section 7.1).

4 Adam

Adam keeps momentum's average direction, RMSProp's average size, and a fix for the zero start of both.

4.1 Two running averages

Adam averages the gradient into m and its square into v, both from zero; this v is not momentum's velocity. At step 1 on the ravine:

m is momentum's velocity scaled by (1 − β₁): an average instead of a sum (SGD, Momentum and Nesterov, section 2).

4.2 The betas set how far back each average looks

4.3 Correcting the zero start

Both averages start at zero, so for a steady gradient the uncorrected step grows.

Dividing each average by the total weight it has received so far makes it a true average from step 1.

m's weight grows about 0.1 per step while v's grows only 0.001, so the raw ratio of m to √v keeps growing; correcting m alone would make step 1 31.6α, which is why both are corrected. Every Adam run takes these corrected first steps, the DQN agent's at 5e-4 too (Deep Q-Networks, section 4).

4.4 One full Adam step

On the first step, x's gradient is −35 and y's is 2. SGD would move x 17.5 times farther than y. Adam moves both by 0.5.

Adam steps by the average direction divided by the average size, times α.

So each coordinate steps about α, whatever the size of its gradient; the divide is element by element (Vectors and Matrices, section 6.3). This is RMSProp's step with m̂ in place of g, both averages corrected. Late in training a coordinate that gets its first gradient steps 3.16α, like uncorrected step 1 above.

The code below is the same five steps.

import numpy as np

def adam_step(theta, g, m, v, t, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8):
    m = b1 * m + (1 - b1) * g
    v = b2 * v + (1 - b2) * g * g
    m_hat = m / (1 - b1 ** t)
    v_hat = v / (1 - b2 ** t)
    theta = theta - lr * m_hat / (np.sqrt(v_hat) + eps)
    return theta, m, v

theta = np.array([-3.5, 5.0])
m = np.zeros(2)
v = np.zeros(2)
for t in range(1, 6):
    g = np.array([10 * theta[0], 0.4 * theta[1]])
    theta, m, v = adam_step(theta, g, m, v, t, lr=0.5)

1. 1The two moving averages, m and v.
2. 2Bias correction of both.
3. 3The step, element by element; returns the new averages too.
4. 4The ravine from (−3.5, 5); the first step is (0.5, −0.5), as in the figure.

Gradient balancing meets this divide in Adam's per-coordinate scaling and the tilt of Adam's step (Gradient Balancing Inside the Training Loop, sections 7 and 7.2).

4.5 How Adam ignores the gradient's scale, and where it stops

Switch a loss from the mean over a batch of 256 examples to the sum, and every gradient is 256 times larger (Gradient Balancing Inside the Training Loop, section 5). Multiply every gradient by c:

Invariance has two limits. The first is ε. For a steady gradient of size |g|:

The second limit is time. Invariance holds for a steady scale, not for one that moves within v's window. After 2,000 steps at gradient 1 the gradient drops to 0.1.

For hundreds of steps Adam steps about 10 times smaller than α, so its step does shrink when the gradient shrinks, just later. LLaMA trains with β₂ 0.95, a window of about 20 steps; with it the same drop recovers to 0.9α 118 steps after the drop.

4.6 What Adam's momentum carries after the gradient changes

Momentum at α 0.01 builds its step up to 10α, about 0.1. Adam at α 0.1 keeps its step near α, also 0.1. Both then coast about 9 last steps. After a turn, √v̂ is still large on the old axis and small on the new one, so the divide boosts the new direction: Adam's step is within 45° of it at once, while m takes 7 steps. That is the case of a multitask run where one task's gradient suddenly dominates a shared weight (Gradient Balancing Inside the Training Loop, section 6.1).

5 Nadam

Nadam puts the Nesterov preview into Adam: it uses the momentum the next step will use.

PyTorch ships it as torch.optim.NAdam, which multiplies β₁ by a warm-up factor, so its momentum weight μ starts near 0.45 and rises toward β₁ over thousands of steps.

6 Racing the optimizers

All four run 200 steps on the ravine; AdaGrad (section 2) and Nadam (Figure 9) are left out to keep four paths readable. The noise toggle adds minibatch-like noise.

With noise on, Adam has the lowest median loss over 20 noise seeds from α 0.07 up, RMSProp at 0.05 and 0.06; no optimizer wins everywhere.

7 Adaptive optimizers in PyTorch

Each runs once per training step. The DQN page's Adam call is this same update.

opt = torch.optim.Adam(model.parameters(), lr=1e-3, betas=(0.9, 0.999), eps=1e-8)
loss.backward()
opt.step()
state = opt.state[next(model.parameters())]
state["exp_avg"]      # m, the first moment
state["exp_avg_sq"]   # v, the second moment
state["step"]         # t, used by the bias correction

1. 1β₁ and β₂ are the betas pair; these are the defaults.
2. 2One step after each backward pass.
3. 3The two state tensors m and v per weight tensor, and the step count t.

8 Failure modes

9 References

1. Duchi, J., et al. (2011). Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research.
2. Tieleman, T., and Hinton, G. (2012). Lecture 6.5, RMSProp. Neural Networks for Machine Learning (Coursera).
3. Pennington, J., Socher, R., and Manning, C. D. (2014). GloVe: global vectors for word representation. EMNLP.
4. Kingma, D. P., and Ba, J. (2015). Adam: a method for stochastic optimization. ICLR.
5. Dozat, T. (2016). Incorporating Nesterov momentum into Adam. ICLR workshop.
6. Touvron, H., et al. (2023). LLaMA: open and efficient foundation language models. arXiv:2302.13971.
7. PyTorch documentation, torch.optim (Adagrad, RMSprop, Adam and NAdam default settings).