AdamW optimizer: what it is, with an example

Adam with weight decay applied separately from the gradient step. It's the default optimizer for training most language models.

What it means

AdamW is the Adam optimizer with weight decay applied separately from the gradient step. Ilya Loshchilov and Frank Hutter proposed it in their paper on decoupled weight decay, and it's now the usual choice for training transformers.

The difference matters because of how Adam works. It rescales each weight's update by the size of its recent gradients. If you add weight decay to the gradient, the old way, Adam rescales the decay too, so it stops acting like decay. AdamW takes the decay out of that machinery.

The difference, on one step

LR, DECAY = 0.01, 0.1

def adam_first_step(grad):
    # On its first step, Adam moves about LR whatever the gradient's size.
    return LR * grad / (abs(grad) + 1e-8)

print(f"{'weight':>8} {'L2 in the gradient':>20} {'AdamW':>10}")
for w in [10.0, 1.0, 0.01]:
    loss_grad = 0.0                                  # say the loss doesn't care about this weight
    with_l2 = w - adam_first_step(loss_grad + DECAY * w)
    adamw = w - adam_first_step(loss_grad) - LR * DECAY * w
    print(f"{w:>8} {with_l2:>20.4f} {adamw:>10.4f}")

What it prints

  weight   L2 in the gradient      AdamW
    10.0               9.9900     9.9900
     1.0               0.9900     0.9990
    0.01               0.0000     0.0100

Here the loss doesn't care about these weights, so only decay should move them. With decay in the gradient, Adam's rescaling takes all three down by the same 0.01, and the smallest weight is wiped out in one step. AdamW shrinks each in proportion to its size, which is what weight decay is meant to do.

In PyTorch

Use torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.1). Passing weight_decay to plain torch.optim.Adam gives you the old behavior, in the gradient, so the two aren't the same thing. You can see AdamW training a real model in build a GPT from scratch.

AdamW

What it means

The difference, on one step

In PyTorch

Try it now

Learn it properly

Other terms from the same lesson

Never trained an AI?