Optimisers
SGD, momentum, and Adam.
Backpropagation tells each weight which way is downhill. But how big a step should it take? Should it remember past steps? The optimiser answers these questions, and a good one can make the difference between a model that trains in an hour, and one that never trains at all.
- Gradients say which way
- How big a step?
- The optimiser decides
Rolling downhill
Plain gradient descent is like walking downhill with no memory, one careful step at a time. In a long narrow valley, it zigzags from wall to wall, wasting time. Momentum is like a heavy ball. It builds up speed in directions that stay consistent, and it smooths out the zigzags. And Adam gives every single weight its own step size, adapting as it goes. That makes it a great default choice.
The optimiser contest
Let's have a contest, on the fashion images, with two passes through the data. Plain S G D, with momentum, Adam, and Adam W, which handles weight decay more correctly. Only one line changes. Run it. Plain S G D falls behind, in this short race. Momentum helps a lot. Adam and Adam W are at the top. With much longer training and careful tuning, S G D can sometimes catch up, but Adam W is what most people start with.
import torch
from torch import nn
from torchvision import datasets, transforms
root = "data"
tf = transforms.ToTensor()
train = datasets.FashionMNIST(root, train=True, transform=tf)
test = datasets.FashionMNIST(root, train=False, transform=tf)
X_test = test.data.float().unsqueeze(1) / 255
def run(make_opt, epochs=2, seed=0, net=None, scheduler=None):
torch.manual_seed(seed)
model = net() if net else nn.Sequential(nn.Flatten(), nn.Linear(784, 256),
nn.ReLU(), nn.Linear(256, 10))
opt = make_opt(model.parameters())
dl = torch.utils.data.DataLoader(train, batch_size=128, shuffle=True)
for _ in range(epochs):
model.train()
for x, y in dl:
loss = nn.functional.cross_entropy(model(x), y)
opt.zero_grad(); loss.backward(); opt.step()
model.eval()
with torch.no_grad():
return (model(X_test).argmax(1) == test.targets).float().mean().item()
from torch.optim import SGD, Adam, AdamW
for name, make in [("SGD", lambda p: SGD(p, lr=0.01)),
("Momentum", lambda p: SGD(p, lr=0.01, momentum=0.9)),
("Adam", lambda p: Adam(p, lr=0.001)),
("AdamW", lambda p: AdamW(p, lr=0.001))]:
print(f"{name:<9} {run(make):.3f}")SGD 0.738 Momentum 0.829 Adam 0.851 AdamW 0.850
Summary
S G D is the simple walker, momentum is the heavy ball, and Adam adapts for every weight. Use Adam W as your default. But one setting matters more than any other. Get it wrong, and nothing works at all.
Key takeaways- Momentum builds speed
- Adam adapts per weight
- AdamW is a strong default