AdamW: steps that adapt

One learning rate cannot suit every weight

Plain gradient descent moves each weight by the learning rate times its gradient. In a network, gradients of different weights can differ by factors of thousands. Suppose one weight's gradient is 100 and another's is 0.01. With a learning rate of 0.001, the first moves by 0.1 and the second by 0.00001. A learning rate small enough for the first weight leaves the second one almost frozen.

The same problem appears in a long, narrow valley of the loss. The steep direction needs small steps. The shallow direction needs large ones.

Try it. Both paths start at the same point on a long, narrow valley, whose lowest point is the cross. Predict first: which path reaches the cross first? Then move the slider to watch them step.

Adam divides the average gradient by its typical size

Adam keeps two numbers for every weight, updated at every step from the new gradient g:

• m, a running average of the gradient: m ← 0.9m + 0.1g;
• v, a running average of the squared gradient: v ← 0.999v + 0.001g².

A running average of this kind keeps most of its old value and mixes in a little of the new one. So it remembers recent values and slowly forgets old ones. The numbers 0.9 and 0.999 are Adam's usual settings, called β₁ and β₂.

Both averages start at 0, which makes them too small for the first steps. Adam corrects this by dividing by 1 − 0.9^t and 1 − 0.999^t at step t. The corrected values are written m̂ and v̂. The step is then:

step = η × m̂ / (√v̂ + ε)

√v̂ is the typical size of this weight's gradient. Dividing by it cancels the gradient's scale. ε is a tiny number, usually 10⁻⁸, that prevents dividing by zero.

The averaging also damps noise. If a weight's gradient keeps flipping sign, m stays near zero and the steps stay small. If the gradient keeps the same sign, the steps stay at about η:

adam_by_hand.py

import numpy as np

# eps is the ε above. 1e-8 means 0.00000001.
def adam_steps(grads, lr=0.001, b1=0.9, b2=0.999, eps=1e-8):
    m = v = 0.0                            # both averages start at 0
    for t, g in enumerate(grads, start=1):  # t counts the steps from 1
        m = b1 * m + (1 - b1) * g          # average gradient
        v = b2 * v + (1 - b2) * g * g      # average squared gradient
        m_hat = m / (1 - b1 ** t)          # correct the start-at-zero bias
        v_hat = v / (1 - b2 ** t)
        print(f"g = {g:6.2f}   step = {lr * m_hat / (np.sqrt(v_hat) + eps):.6f}")

adam_steps([5.0, 5.0, 5.0, 5.0])      # steady gradient
print()
adam_steps([5.0, -5.0, 5.0, -5.0])    # gradient that keeps flipping

g =   5.00   step = 0.001000
g =   5.00   step = 0.001000
g =   5.00   step = 0.001000
g =   5.00   step = 0.001000

g =   5.00   step = 0.001000
g =  -5.00   step = -0.000053
g =   5.00   step = 0.000336
g =  -5.00   step = -0.000053

The printed step is the amount subtracted from the weight. PyTorch's Adam gives the same first step for the worked example:

adam_torch.py

import torch

# Two weights whose gradients differ in size by a factor of 10,000.
w = torch.tensor([1.0, 1.0], requires_grad=True)
opt = torch.optim.Adam([w], lr=0.001)   # an Adam optimizer that updates w

w.grad = torch.tensor([100.0, 0.01])    # set the gradients by hand
opt.step()                              # take one Adam step
print(1.0 - w.detach())   # how far each weight moved

tensor([0.0010, 0.0010])

torch.tensor makes a PyTorch array, called a tensor. With requires_grad=True, the tensor keeps a gradient in .grad, as Automatic gradients in PyTorch explains. w.detach() gives the same values without gradient tracking, so the output holds only the values.

Weight decay shrinks every weight a little at every step

Weight decay multiplies every weight by a number just below 1 at each step. A weight that the loss does not need drifts toward zero. This keeps weights small, which tends to help a model work on new examples, not only on the ones it trained on. The strength is a number λ (lambda), often between 0.01 and 0.1.

AdamW applies the decay beside the Adam step, not through it. The W in the name stands for this "decoupled" weight decay. One AdamW step is:

w_new = w × (1 − ηλ) − η × m̂ / (√v̂ + ε)

The older way added λw to the gradient instead, before the averages. For plain gradient descent the two ways are the same. For Adam they are not. The added term would also be divided by √v̂, so weights with large gradients would hardly shrink, and weights with small gradients would shrink a lot. Loshchilov and Hutter proposed decoupling the decay in a paper published in 2019. In their tests it gave better results than Adam with the older decay, and AdamW has since become the usual choice.

adamw_decay.py

import torch

w = torch.tensor([2.0], requires_grad=True)
opt = torch.optim.AdamW([w], lr=0.01, weight_decay=0.1)

w.grad = torch.zeros(1)   # no pull from the loss at all
opt.step()
print(w.item())           # 2 * (1 - 0.01 * 0.1) = 1.998

1.9980000257492065

.item() turns the one-number tensor into a plain Python number. The printed value is 1.998 up to the rounding of 32-bit numbers. PyTorch's AdamW defaults are lr=0.001, betas=(0.9, 0.999), eps=1e-08 and weight_decay=0.01.