Optimizers

Adam: momentum plus a step size per parameter

Momentum fixes the bouncing. A second problem remains. In a network, some weights get big gradients and others get tiny ones. With one learning rate for all of them, the big-gradient weights take huge steps. The tiny-gradient weights barely move.

Think of the volume-leveling setting in a music app. Loud songs get turned down and quiet songs get turned up, so everything plays at about the same volume. Adam does this for gradients. It tracks how big each weight's gradients usually are, and divides by that. Big gradients get turned down, and small ones get turned up. Adam also uses momentum. Its name is short for "adaptive moment estimation".

Two running averages

Picture tracking your weekly grocery bill with a running average. Each week, you keep 90% of the old average and mix in 10% of this week's bill. Recent weeks count most, and old weeks slowly fade away.

Adam keeps two averages like that for every weight. The first, $m$, is a running average of the gradient. Its keep-fraction is $\beta_1$ (beta one), usually 0.9.

In plain words: keep 90% of the old average, and add 10% of this step's gradient. It's momentum again, in a slightly different form.

The second average uses the squared gradient, $g^2$. Squaring makes every gradient positive, so this average ignores direction. It only tracks how big the gradients usually are. Adam calls it $v$. (That's a new meaning for v. In Adam, it isn't a velocity.) Its keep-fraction is $\beta_2$ (beta two), usually 0.999.

In plain words: keep 99.9% of the old average of squared gradients, and add 0.1% of this step's squared gradient.

Both averages start at 0. Say the first gradient is 2. Then m = 0.9 × 0 + 0.1 × 2 = 0.2. And v = 0.999 × 0 + 0.001 × 4 = 0.004.

b1, b2 = 0.9, 0.999
g = 2.0                              # this step's gradient
m = b1 * 0.0 + (1 - b1) * g          # average of gradients, starting from 0
v = b2 * 0.0 + (1 - b2) * g ** 2     # average of squared gradients, starting from 0
print(round(m, 4), round(v, 4))

It prints 0.2 and 0.004. Both are far too small, because starting at 0 drags them down. Adam fixes this with a bias correction, written $\hat m$ and $\hat v$ ("m-hat" and "v-hat"). On this first step, it turns 0.2 back into 2 and 0.004 back into 4. Its effect fades as training goes on.

t = 1
print(round(0.2 / (1 - 0.9 ** t), 4), round(0.004 / (1 - 0.999 ** t), 4))

The update

One more symbol: $\epsilon$ (epsilon) is a tiny number, such as 0.00000001. It's there only so Adam never divides by zero.

In plain words: divide the averaged gradient by the square root of the averaged squared gradient. Multiply by the learning rate, and step downhill.

The square root undoes the squaring, so the top and bottom of that fraction have similar sizes. When a weight's gradients are steady, the fraction comes out close to 1, however big or small the gradients are. In the example above, it's 2 ÷ √4 = 2 ÷ 2 = 1. When gradients keep flipping direction, $\hat m$ averages out small, and the step shrinks.

So weights with large, steady gradients don't take huge steps, and weights with tiny gradients aren't left behind. Each gets a step of roughly size $\eta$. That makes Adam much less sensitive to the learning rate. Its defaults are $\beta_1 = 0.9$, $\beta_2 = 0.999$ and $\eta = 10^{-3}$, which is 0.001.

import math

lr = 0.01                          # learning rate
b1, b2, eps = 0.9, 0.999, 1e-8     # Adam's defaults

for g in [100.0, 0.001]:           # a loud gradient and a quiet one
    m = (1 - b1) * g               # average of gradients (it started at 0)
    v = (1 - b2) * g ** 2          # average of squared gradients (it started at 0)
    m_hat = m / (1 - b1)           # bias correction for step 1
    v_hat = v / (1 - b2)
    adam = lr * m_hat / (math.sqrt(v_hat) + eps)
    sgd = lr * g                   # plain gradient descent, for comparison
    print(f"gradient {g}: plain SGD moves {sgd:.5f}, Adam moves {adam:.4f}")

Three optimizers in one valley

Now watch all three on the canyon from the start, $f(x, y) = x^2 + 50 y^2$, from the same starting point. As you saw, plain gradient descent's learning rate can't go above 0.02 here. Any bigger, and the steep y direction blows up.

The code uses PyTorch's built-in optimizers. torch.optim.SGD is plain gradient descent, and adding momentum=0.9 turns it into the momentum rule above. Adam gets a learning rate of 0.3: since it moves each weight about one learning rate per step, that number sets its stride directly.

The plot draws each path over gray contour lines. They work like the height lines on a hiking map: each ring joins spots with the same loss.

import numpy as np
import torch
import matplotlib.pyplot as plt

def path(make_opt, steps=60):
    p = torch.tensor([-4.0, 1.0], requires_grad=True)
    opt = make_opt([p])
    points = [p.detach().numpy().copy()]
    for _ in range(steps):
        opt.zero_grad()
        loss = p[0] ** 2 + 50 * p[1] ** 2
        loss.backward()
        opt.step()
        points.append(p.detach().numpy().copy())
    return np.array(points), loss.item()

runs = {"SGD, lr 0.019": lambda ps: torch.optim.SGD(ps, lr=0.019),
        "SGD + momentum 0.9, lr 0.004": lambda ps: torch.optim.SGD(ps, lr=0.004, momentum=0.9),
        "Adam, lr 0.3": lambda ps: torch.optim.Adam(ps, lr=0.3)}
gx, gy = np.meshgrid(np.linspace(-5, 5, 200), np.linspace(-1.3, 1.3, 200))
plt.contour(gx, gy, gx ** 2 + 50 * gy ** 2, levels=15, colors="gray", linewidths=0.5)
for name, make in runs.items():
    pts, final = path(make)
    plt.plot(pts[:, 0], pts[:, 1], marker=".", markersize=3, label=name)
    print(f"{name:<29} loss after 60 steps: {final:.4f}")
plt.legend()

Plain SGD zig-zags across the valley. After 60 steps, it's still far from the lowest point. Momentum builds up speed along the valley floor. Adam takes roughly equal-sized steps in both directions from the start. Both end several times closer to the lowest point than plain SGD.

On a real network

Here's the digit classifier again, trained for 10 epochs with batches of 32. Each optimizer gets a sensible learning rate. The code prints each one's loss and accuracy on the validation set, and plots the validation loss after every epoch.

import torch
import torch.nn as nn
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader, TensorDataset
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

digits = load_digits()
X = torch.tensor(digits.data / 16.0, dtype=torch.float32)
y = torch.tensor(digits.target)
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=0)
loss_fn = nn.CrossEntropyLoss()

choices = {"SGD (lr 0.05)": lambda ps: torch.optim.SGD(ps, lr=0.05),
           "SGD + momentum (lr 0.05)": lambda ps: torch.optim.SGD(ps, lr=0.05, momentum=0.9),
           "Adam (lr 0.001)": lambda ps: torch.optim.Adam(ps, lr=0.001)}
for name, make in choices.items():
    torch.manual_seed(0)
    model = nn.Sequential(nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 10))
    opt = make(model.parameters())
    loader = DataLoader(TensorDataset(X_train, y_train), batch_size=32, shuffle=True)
    curve = []
    for epoch in range(10):
        for xb, yb in loader:
            opt.zero_grad()
            loss_fn(model(xb), yb).backward()
            opt.step()
        with torch.no_grad():
            curve.append(loss_fn(model(X_val), y_val).item())
            acc = (model(X_val).argmax(1) == y_val).float().mean().item()
    print(f"{name:<26} val loss {curve[-1]:.3f}  val accuracy {acc:.3f}")
    plt.plot(range(1, 11), curve, marker="o", label=name)
plt.xlabel("epoch")
plt.ylabel("validation loss")
plt.legend()

Both momentum and Adam beat plain SGD here. With these settings, momentum edges out Adam. That's a reminder that which optimizer "wins" depends on tuning. Adam's advantage is that its default settings work reasonably well almost everywhere.

AdamW and weight decay

Weight decay shrinks every weight a little on every step. Weights the data doesn't keep pushing on drift toward zero, which helps the model do well on new data. It's the L2 idea from Phase 3's regularization lesson, which penalized big weights.

There are two ways to add weight decay to Adam. The old way adds it to the gradient. But then Adam's volume leveling rescales it along with everything else. Weights with big gradients end up with less decay than you asked for. That interacts badly with the per-parameter scaling.

AdamW applies the decay directly to the weights instead, outside the leveling. This is called "decoupled" weight decay. It works better, and it's the standard choice for transformers, the model design behind LLMs. In PyTorch, each AdamW step multiplies every weight by 1 − learning rate × weight decay, on top of the usual Adam step.

A typical LLM setting is:

torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.1)

Here 3e-4 means 3 × 10⁻⁴, or 0.0003. So each step multiplies every weight by 1 − 0.0003 × 0.1 = 0.99997. That's a tiny shrink, but training runs take many thousands of steps.

The memory bill

Better optimizers need memory to remember things:

• Plain SGD needs nothing extra.
• Momentum stores one velocity per parameter.
• Adam stores two numbers per parameter, $m$ and $v$.

That adds up fast. A 32-bit number takes 4 bytes. So a model with 7 billion parameters needs 28 GB for its weights alone: 7 billion × 4 bytes. Adam's state is twice that, 56 GB. And that's before the gradients and activations (the in-between results saved for the backward pass).

Optimizer memory is a big part of why training needs far more GPU memory than running a model. It's also why memory-saving optimizers exist, and low-precision tricks that store each number in fewer bytes.

Rules of thumb

• Start with AdamW (or Adam). Use a learning rate of $10^{-3}$ for small networks, and $3 \times 10^{-4}$ or lower for transformers. If the loss explodes, lower it.
• SGD with momentum is still competitive for some vision models, and can generalize slightly better. But it needs more learning-rate tuning.
• Pair any optimizer with a schedule: a warmup, then a decay. You met these in the learning-rate lesson in Phase 3.