Optimizers: momentum and Adam
Every lesson so far trained with plain SGD: step each parameter by the learning rate times its gradient.
Almost every modern model, ACT included, trains with Adam instead. What does Adam do differently, and is it worth it?
Overview
An optimizer is the rule that turns gradients into parameter changes. optimizer.step() runs it.
The mental model: plain SGD looks only at today’s gradient; momentum and Adam also remember earlier ones. Momentum remembers their direction, Adam also their size.
The valley
To see the difference, we need a loss where plain SGD struggles. Here is one with two parameters:
\[L(w_1, w_2) = \tfrac12\left(w_1^2 + 100\,w_2^2\right)\]Its lowest point is (0, 0). Across the valley, along \(w_2\), it is 100 times steeper than along it.
Real networks have such valleys too: some parameters change the loss far more than others. The T-shirt network in the several-inputs lesson scaled its inputs to avoid one, “so one learning rate suits both”.
The gradient is \((w_1,\ 100\,w_2)\). Every optimizer below starts at \((-4, 1)\), where the gradient is \((-4, 100)\), and runs until the loss is below 0.001.
Valley lab
How to read the map:
- Each oval joins points of equal loss; the bottom is the cross at (0, 0).
- The bold path is the selected optimizer, one dot per step; the other two are drawn faint. The legend on the map names each colour: orange plain SGD, magenta momentum, blue Adam.
- Each optimizer runs at its own fastest rate, measured in the table below.
Plain SGD
Plain SGD subtracts the rate times the gradient. One step at rate , worked by hand:
\[(-4,\ 1) - 0.019 \times (-4,\ 100) = (-3.924,\ -0.9)\]Each parameter is multiplied by its own factor every step:
- Across the valley, \(w_2\) is multiplied by \(1 - 100 \times 0.019 = -0.9\). It jumps to the other side, a little closer.
- Along the valley, \(w_1\) is multiplied by \(1 - 0.019 = 0.981\). It gains only 2% per step.
A bigger rate would speed up \(w_1\). But above 0.02 the factor across falls below −1, and every jump lands farther out than the last. So the steep direction caps the rate, and the gentle direction pays for it: steps.
The number of steps to the bottom, for every rate tried:
A dash means the loss never got below 0.001 in 300 steps: the rate was too small to arrive, or so large the steps grew.
Momentum
Momentum keeps a running sum of past gradients, the velocity, and steps along that instead of along today’s gradient:
\[\text{velocity} \leftarrow 0.9 \times \text{velocity} + \text{gradient}\qquad w \leftarrow w - \text{rate} \times \text{velocity}\]def momentum_step(w, g, velocity, lr, beta=0.9):
"""velocity = beta × velocity + gradient; then step along the velocity (PyTorch's form)."""
velocity = [beta * v + gi for v, gi in zip(velocity, g)] if velocity else list(g)
return [wi - lr * v for wi, v in zip(w, velocity)], velocityFull file · pytorch_optimizers.py · 152 lines
"""Why almost every modern model trains with Adam: plain SGD, momentum, Adam and AdamW, first on a long narrow
valley in two parameters, then on the bouncing-ball network.
Run: python3 site/ml/foundations/examples/pytorch_optimizers.py
"""
import math
import sys
from pathlib import Path
import numpy as np
import torch
from torch import nn
# ---- the valley: a loss over two parameters, 100 times steeper across than along ---------------------------
START = (-4.0, 1.0)
def valley_loss(w):
return 0.5 * (w[0] ** 2 + 100 * w[1] ** 2)
OPTIMIZERS = {
"SGD": lambda params, lr: torch.optim.SGD(params, lr=lr),
"momentum": lambda params, lr: torch.optim.SGD(params, lr=lr, momentum=0.9),
"Adam": lambda params, lr: torch.optim.Adam(params, lr=lr),
"AdamW": lambda params, lr: torch.optim.AdamW(params, lr=lr, weight_decay=0.01),
}
def descend(name, lr, steps=300, target=1e-3):
"""Run one optimizer from START. Returns the path and the first step whose loss is below `target` (or None)."""
w = torch.tensor(START, dtype=torch.float64, requires_grad=True)
optimizer = OPTIMIZERS[name]([w], lr)
path, reached = [w.tolist()], None
for step in range(steps):
loss = valley_loss(w)
if not math.isfinite(loss.item()) or loss.item() > 1e6:
break # the steps grew instead of shrinking
if reached is None and loss.item() < target:
reached = step
optimizer.zero_grad()
loss.backward()
optimizer.step()
path.append(w.tolist())
return path, reached
# ---- the same updates written by hand, one parameter list at a time ---------------------------------------
def momentum_step(w, g, velocity, lr, beta=0.9):
"""velocity = beta × velocity + gradient; then step along the velocity (PyTorch's form)."""
velocity = [beta * v + gi for v, gi in zip(velocity, g)] if velocity else list(g)
return [wi - lr * v for wi, v in zip(w, velocity)], velocity
def adam_step(w, g, m, v, t, lr, beta1=0.9, beta2=0.999, eps=1e-8, decoupled_decay=0.0):
"""m: running average of gradients. v: running average of squared gradients. t: the step number, from 1."""
m = [beta1 * mi + (1 - beta1) * gi for mi, gi in zip(m, g)]
v = [beta2 * vi + (1 - beta2) * gi * gi for vi, gi in zip(v, g)]
m_hat = [mi / (1 - beta1 ** t) for mi in m] # undo the pull toward the zero start
v_hat = [vi / (1 - beta2 ** t) for vi in v]
w = [wi * (1 - lr * decoupled_decay) for wi in w] # AdamW: shrink every weight by the same fraction
return [wi - lr * mh / (math.sqrt(vh) + eps) for wi, mh, vh in zip(w, m_hat, v_hat)], m, v
# ---- the real check: the bouncing-ball network, 4 → 16 → 4, with only the optimizer line changed ----------
LAB = Path(__file__).resolve().parents[4] / "classes/ml/labs/bouncing-ball"
def ball_data():
data = np.load(LAB / "data/throws.npz")
train = data["states"][data["split"] == "train"]
test = data["states"][data["split"] == "test"]
contact = data["contact"][data["split"] == "test"]
pairs = lambda t: (torch.tensor(t[:, :-1].reshape(-1, 4), dtype=torch.float32),
torch.tensor((t[:, 1:] - t[:, :-1]).reshape(-1, 4), dtype=torch.float32))
# Whole throws: 450 to train on, the last 50 training throws to choose the learning rate, 100 to test once.
return pairs(train[:450]), pairs(train[450:]), pairs(test), contact.reshape(-1)
def train_ball(name, lr, seed=0, epochs=240, batch=256):
sys.path.insert(0, str(LAB))
from model import make_network
torch.set_num_threads(1)
torch.manual_seed(seed)
(inputs, targets), validation, test, contact = ball_data()
network = make_network(16)
optimizer = OPTIMIZERS[name](network.parameters(), lr) # the only line that differs between the four runs
loss_fn = nn.MSELoss()
order = torch.Generator().manual_seed(seed)
for epoch in range(epochs):
for rows in torch.randperm(len(inputs), generator=order).split(batch):
loss = loss_fn(network(inputs[rows]), targets[rows])
optimizer.zero_grad()
loss.backward()
optimizer.step()
with torch.no_grad():
validation_loss = loss_fn(network(validation[0]), validation[1]).item()
miss = network(test[0]) - test[1]
air = torch.linalg.norm(miss[~contact, :2], dim=1).pow(2).mean().sqrt().item() * 1000 # mm, in flight
bounce = torch.linalg.norm(miss[contact, 2:], dim=1).pow(2).mean().sqrt().item() # m/s, at the floor
return {"validation_loss": validation_loss, "miss_air_mm": air, "miss_bounce": bounce}
def tshirt_check(name, lr, seed=0, steps=3000, target=0.1):
"""The T-shirt network from Neural network: several inputs (scaled inputs, 80 customers): the first step whose training
loss is below `target` (None if never), and how many of 200 new customers it sizes right at the end."""
from pytorch_choices import make_people, make_network, scale
torch.set_num_threads(1)
torch.manual_seed(seed)
people, sizes = make_people(80, torch.Generator().manual_seed(seed))
network = make_network()
optimizer = OPTIMIZERS[name](network.parameters(), lr)
loss_fn, reached = nn.CrossEntropyLoss(), None
for step in range(steps):
loss = loss_fn(network(scale(people)), sizes)
if reached is None and loss.item() < target:
reached = step
optimizer.zero_grad()
loss.backward()
optimizer.step()
new_people, new_sizes = make_people(200, torch.Generator().manual_seed(99))
with torch.no_grad():
right = (network(scale(new_people)).argmax(dim=1) == new_sizes).sum().item()
return {"reached": reached, "right": right}
if __name__ == "__main__":
print("The valley, from (-4, 1): first step with loss below 0.001")
for name, lr in [("SGD", 0.019), ("momentum", 0.015), ("Adam", 0.3)]:
print(f" {name:8s} lr {lr}: step {descend(name, lr)[1]}")
g = [-4.0, 100.0] # the gradient at the start: (w1, 100 × w2)
w, velocity = momentum_step(list(START), g, None, 0.015)
print("momentum, step 1 by hand:", [round(x, 4) for x in w])
path, _ = descend("momentum", 0.015, steps=1)
print("momentum, step 1 by torch.optim.SGD(momentum=0.9):", [round(x, 4) for x in path[1]])
w, m, v = adam_step(list(START), g, [0, 0], [0, 0], 1, 0.3)
print("Adam, step 1 by hand:", [round(x, 4) for x in w])
print("Adam, step 1 by torch.optim.Adam:", [round(x, 4) for x in descend("Adam", 0.3, steps=1)[0][1]])
# 50 hand-written Adam steps against PyTorch's.
w, m, v = list(START), [0.0, 0.0], [0.0, 0.0]
for t in range(1, 51):
w, m, v = adam_step(w, [w[0], 100 * w[1]], m, v, t, 0.3)
print("Adam after 50 steps, by hand:", [f"{x:.6f}" for x in w], " torch:", [f"{x:.6f}" for x in descend("Adam", 0.3, steps=50)[0][50]])
print("\nThe ball (seed 0, 240 epochs): validation loss and test misses")
for name, lr in [("SGD", 0.1), ("Adam", 0.001)]: # the rates chosen on the validation throws
r = train_ball(name, lr)
print(f" {name:5s} lr {lr}: validation {r['validation_loss']:.5f}, miss in the air {r['miss_air_mm']:.2f} mm, "
f"at the bounce {r['miss_bounce']:.2f} m/s")
The first three steps at rate :
- Along the valley the gradient always points the same way, so the velocity keeps growing: −4, −7.54, −10.61. It can reach 10 times one gradient, since \(1 + 0.9 + 0.9^2 + \dots = 10\).
- Across it the velocity turns round each time \(w_2\) crosses the floor, but late. Instead of a zig-zag at every step, the path swings in wide arcs that slowly die down.
The speed along the floor wins: steps instead of , even though the path overshoots the bottom and swings back. PyTorch’s torch.optim.SGD(..., momentum=0.9) gives the same numbers.
- Momentum
- stepping along a running sum of past gradients instead of the current one.
Adam
Adam gives every parameter its own step size. It divides each parameter’s gradient by the typical size of that parameter’s recent gradients, so a steep parameter and a gentle one move at a similar pace.
It keeps two running averages per parameter:
- m, the average gradient: momentum, as above.
- v, the average squared gradient: how big this parameter’s gradients have been.
def adam_step(w, g, m, v, t, lr, beta1=0.9, beta2=0.999, eps=1e-8, decoupled_decay=0.0):
"""m: running average of gradients. v: running average of squared gradients. t: the step number, from 1."""
m = [beta1 * mi + (1 - beta1) * gi for mi, gi in zip(m, g)]
v = [beta2 * vi + (1 - beta2) * gi * gi for vi, gi in zip(v, g)]
m_hat = [mi / (1 - beta1 ** t) for mi in m] # undo the pull toward the zero start
v_hat = [vi / (1 - beta2 ** t) for vi in v]
w = [wi * (1 - lr * decoupled_decay) for wi in w] # AdamW: shrink every weight by the same fraction
return [wi - lr * mh / (math.sqrt(vh) + eps) for wi, mh, vh in zip(w, m_hat, v_hat)], m, vbeta1andbeta2set how long each average remembers: about 10 steps for m, about 1,000 for v.m_hatandv_hatcorrect for m and v starting at zero; without it the first steps would be far too small.epsonly prevents a division by zero.
The first step, by hand at rate . After the correction, \(\hat m\) is the gradient itself and \(\sqrt{\hat v}\) is its size, so each parameter moves by exactly the rate:
\[\frac{-4}{4} = -1,\quad \frac{100}{100} = 1 \quad\Rightarrow\quad (-4,\ 1) - 0.3 \times (-1,\ 1) = (-3.7,\ 0.7)\]The gentle \(w_1\) and the steep \(w_2\) move by the same 0.3. The first three steps:
After 50 steps the hand-written Adam is at , and torch.optim.Adam at : identical to six decimals.
On the valley Adam needs steps. More important is the rate table above: Adam reaches the bottom at every rate from 0.03 to 1.0, while plain SGD needs a rate between 0.015 and 0.019. That forgiveness is much of why Adam became the default.
- Adam
- momentum plus a per-parameter step size: each gradient is divided by the root of its recent average square.
The ball
A two-parameter valley is a drawing. The real question is whether Adam helps a network you know: the bouncing-ball workshop’s 4 → 16 → 4, which learned with plain SGD at rate 0.1.
network = make_network(16)
optimizer = OPTIMIZERS[name](network.parameters(), lr) # the only line that differs between the four runsThe comparison follows Generalization and evaluation:
- Training uses 450 of the 500 training throws, 240 epochs, batches of 256, as in the workshop.
- Each optimizer’s rate is chosen from three or four on the other 50 training throws, by median validation loss over seeds 0–4.
- The 100 test throws grade the chosen rates once, over seeds 0–9.
The two grades are the workshop’s: the miss in the air, in millimetres, and the miss at the bounce, in metres per second.
Every seed and every rate tried
Adam does not always win. On the T-shirt network from the several-inputs lesson, small and with scaled inputs, the same comparison goes the other way. Steps until the training loss falls below 0.1, median of seeds 0–9, and new customers sized right out of 200:
Weight decay
Weight decay shrinks every weight a little toward zero at each step. It keeps weights from growing large without need, one more guard against overfitting.
Plain SGD adds \(\lambda w\) to each gradient, and that is the same as shrinking every weight by the fraction \(\text{rate} \times \lambda\). Adam breaks that equivalence, because it divides each gradient by its own size.
One step from weights (1, 1), at rate 0.1 and weight decay 0.01. Weight 1 gets no gradient, weight 2 a gradient of 10:
Adam(weight_decay=0.01)adds the decay to the gradient, then divides by its size. For weight 1, the tiny decay term alone becomes a full step: it shrinks by 10% at once.AdamW(weight_decay=0.01)shrinks every weight by the same \(0.1 \times 0.01 = 0.1\%\), separately from the gradient step.
- Weight decay
- shrinking every weight toward zero by a small fraction each step.
- AdamW
- Adam with weight decay applied separately (decoupled), so every weight shrinks by the same fraction.
This is the optimizer ACT uses. In LeRobot 0.6.1 its preset is AdamW with:
- rate \(10^{-5}\), for the whole network including its image backbone;
- weight decay \(10^{-4}\);
- betas (0.9, 0.999) and eps \(10^{-8}\), PyTorch’s defaults, the same as above;
- gradients clipped to a total size of at most 10 before each step.
Schedules
A learning-rate schedule changes the rate during training. Two are common:
- Warm-up starts with a tiny rate and raises it over the first steps, while Adam’s averages are still unreliable.
- Decay lowers the rate toward the end, so the last steps settle instead of bouncing around the bottom.
Large language models nearly always use both. ACT in LeRobot uses neither: its preset has no scheduler, so the rate stays at \(10^{-5}\) for the whole run.
- Learning-rate schedule
- a rule that changes the learning rate as training goes on.
Recap
valley_loss(w)is a loss 100 times steeper across than along: plain SGD’s rate is capped by the steep direction.momentum_step(...)steps along a running sum of gradients: it builds speed along the floor and cancels the zig-zag.adam_step(...)adds a running average of squared gradients and divides by its root: every parameter gets its own step size.descend(name, lr)runs the same updates withtorch.optim; the hand-written steps match it.train_ball(name, lr, seed)is the workshop’s training loop with one line changed: the optimizer.torch.optim.AdamW(params, lr, weight_decay)is what ACT trains with: Adam plus decay applied separately.
From now on the lessons may use Adam or AdamW. Each will say which, and with what rate.
Key concepts
- Optimizer
- The rule that turns gradients into parameter changes;
optimizer.step(). - Plain SGD
- Parameter minus rate times gradient.
- Momentum
- Stepping along a running sum of past gradients.
- Adam
- Momentum plus a per-parameter step size from the average squared gradient.
- Bias correction
- Adam’s division by \(1 - \beta^t\), which undoes the averages’ zero start.
- Weight decay
- Shrinking every weight toward zero by a small fraction each step.
- AdamW
- Adam with weight decay applied separately from the gradient step.
- Learning-rate schedule
- Changing the rate during training: warm-up, decay.
- Gradient clipping
- Scaling the gradients down when their total size exceeds a limit.