Adam — Module 6.1 | Codextroop

How This Lesson Fits the Module

SGD uses one global learning rate for all parameters. Adam (Adaptive Moment Estimation) scales each parameter’s step using running estimates of the gradient mean and variance.

Adam is the default starting point for many practitioners. Understanding its mechanics helps you debug divergence and choose AdamW for transformers.

Learning Objectives

By the end of this lesson, students should be able to:

Describe Adam’s first and second moment estimates (m and v).
Configure torch.optim.Adam and AdamW in PyTorch.
Explain bias correction in early training steps.
Contrast Adam with SGD on convergence speed and final generalization.
Know sensible defaults: lr=1e-3, betas=(0.9, 0.999), eps=1e-8.

How Adam Works

At each step, Adam maintains:

m_t — exponential moving average of gradients (momentum-like).
v_t — exponential moving average of squared gradients (adaptive scale).

The update divides by √v_t + ε, giving larger steps to parameters with small recent gradients.

Definition — Adam

Adam combines ideas from RMSprop and momentum: per-parameter adaptive learning rates from second moments, plus momentum from first moments. AdamW decouples weight decay from the adaptive step, fixing a regularization bug in the original Adam formulation.

Hyperparameter	Default	Role
`lr`	1e-3	Global step scale
`betas`	(0.9, 0.999)	Decay for m and v
`eps`	1e-8	Numerical stability in denominator
`weight_decay`	0 (use AdamW)	L2-style shrinkage

PyTorch Adam vs AdamW

import torch.nn as nn

model = nn.TransformerEncoderLayer(d_model=128, nhead=4, batch_first=True)

# Prefer AdamW for transformers and modern DL
optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=3e-4,
    betas=(0.9, 0.98),
    weight_decay=0.01,
)

# Classic Adam — still common in tutorials
optimizer_adam = torch.optim.Adam(model.parameters(), lr=1e-3)

# Training loop unchanged: zero_grad → backward → step

Practical Default Start with AdamW(lr=1e-3, weight_decay=0.01) on tabular and NLP tasks; switch to SGD+momentum only when benchmarks demand it.

Common Misconception: “Adam eliminates the need to tune learning rate.”

Reality: Adam is less sensitive than SGD but wrong lr still diverges or stalls. Use learning rate finders or a short grid search.

Common Misconception: “Adam and AdamW are identical.”

Reality: AdamW applies weight decay directly to weights, not through the gradient-adapted step. For regularized training, prefer AdamW.

Critical Mistake — Adam + Very Large Batch

Linear scaling rules designed for SGD do not always transfer. Very large batch sizes with Adam may need lr warmup and careful validation—do not blindly 10× the learning rate.

Knowledge Check

Short Answer: What two moments does Adam track? Answer: Mean of gradients (m) and mean of squared gradients (v).
True/False: Adam uses the same step size for every parameter. Answer: False—it adapts per parameter.
Multiple Choice: AdamW improves: (a) batch norm, (b) decoupled weight decay, (c) dropout, (d) softmax. Answer: (b).
Short Answer: Default Adam learning rate in PyTorch? Answer: 1e-3.
True/False: eps prevents division by zero in the adaptive denominator. Answer: True.
Multiple Choice: betas=(0.9, 0.999) control decay of: (a) batch size, (b) m and v, (c) epochs, (d) dropout rate. Answer: (b).
Short Answer: Why bias correction in early steps? Answer: m and v are initialized at zero and biased toward zero initially.
True/False: Adam always generalizes better than SGD. Answer: False—task dependent.
Multiple Choice: Next topic on data grouping: (a) Batch, (b) Residual Networks, (c) Tanh, (d) Bias. Answer: (a).
Short Answer: Three calls after forward+loss in training loop. Answer: zero_grad, backward, step.

Key Takeaways

Adam adapts step sizes using gradient first and second moments.
AdamW is the preferred variant when using weight decay.
Defaults work surprisingly often—still validate lr on your task.
Adam converges fast early; SGD may win at the end on some vision jobs.
Next: Batch — how data is grouped for each update.

Trainer’s Guide

Demo: Plot per-parameter update magnitudes for Adam vs SGD on the first 100 steps of a small MLP.

Reading: Point students to the original Adam paper—focus on bias correction, not the full proof.

What’s Next Optimizers consume batches of data each step.