SGD with a fixed learning rate is the baseline everyone learns first. It’s also almost never what I run for transformers. I spent my first month of GPT training tweaking learning rates manually, watching loss spike when a batch had weird gradients, then discovering that Adam — specifically AdamW — does most of that annoyance management for me.
This post is the optimizer companion to cross-entropy loss. Same deal as autograd: I’ll build the update rule from SGD upward, derive bias correction, show exactly where AdamW diverges from Adam, and run one scalar parameter through three timesteps with real numbers. Defaults match what I use in AL-1: lr=3e-4, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.1.
Why isn’t plain SGD enough?
Stochastic gradient descent:
where $g_t = \nabla_\theta L$ at step $t$ and $\eta$ is learning rate.
Two problems show up immediately in deep networks:
- Different scales per parameter. Embedding gradients and LayerNorm gradients can differ by orders of magnitude. One global $\eta$ is always wrong for someone.
- Noisy minibatches. A single batch gradient is a rough estimate. Raw SGD jitters, especially early in training.
Momentum and adaptive methods fix these. Adam combines both.
flowchart TB
SGD["SGD\nfixed η"] --> MOM["+ Momentum\nsmooth direction"]
MOM --> RMS["+ RMSprop\nper-param scale"]
RMS --> ADAM["Adam\nbias correction"]
ADAM --> ADAMW["AdamW\ndecoupled WD"]
What does momentum add?
Exponential moving average of gradients:
With $\beta_1 = 0.9$, the update is a smoothed blend of recent gradients — dampens oscillation in ravines, builds speed along consistent directions.
Initialize $m_0 = 0$. At $t = 1$:
The first step only uses 10% of the actual gradient. That’s the cold start problem Adam fixes with bias correction.
What does RMSprop add?
Track a moving average of squared gradients:
(elementwise square)
Scale each parameter’s update inversely:
| Large historical $ | g | $ → large $v$ → smaller effective step. Small $v$ → bigger step. That’s the per-parameter adaptive learning rate. |
$\epsilon = 10^{-8}$ prevents division by zero. In float32 training I sometimes use $10^{-6}$; for my NumPy AL-1 build $10^{-8}$ is fine.
How does Adam combine them?
Adam (Kingma & Ba): momentum on $g$, RMSprop on $g^2$, plus bias correction.
At step $t$ (1-indexed):
Bias correction:
Update:
Why bias correction? Both $m_t$ and $v_t$ start at zero, so early $m_t$ and $v_t$ are systematically underestimated. Dividing by $1 - \beta^t$ ramps the effective average up toward the true scale. At $t = 1$, $\beta_1 = 0.9$:
First-step gradient gets full weight instead of 10%.
At $t = 10$: $1 - 0.9^{10} \approx 0.651$. Still a noticeable boost. By $t = 100$ correction is negligible for $\beta_1$, but $\beta_2 = 0.999$ needs hundreds of steps to settle — $v_t$ warms up slowly because squared gradients are tiny early on.
A numeric walk-through: one weight, three steps
Single parameter $\theta$, learning rate $\eta = 0.1$, $\beta_1 = 0.9$, $\beta_2 = 0.999$, $\epsilon = 10^{-8}$, no weight decay.
Start: $\theta_0 = 1.0$, $m_0 = 0$, $v_0 = 0$.
Step 1: $g_1 = 0.5$
Without bias correction I’d have used $m_1 = 0.05$, step size $0.01$ — ten times too small on step one.
Step 2: $g_2 = -0.3$
Gradient flipped sign; momentum ($m$ still positive) fights the flip briefly — that’s the smoothing.
Step 3: $g_3 = 0.2$ — I’ll stop writing every digit. Point is: $v_t$ grows when gradients are consistently large, shrinking effective steps on noisy axes.
Adam vs AdamW: what actually changes?
L2 regularization in Adam (the old way): add $\lambda \theta$ to the gradient before the Adam update:
Weight decay gets entangled with the adaptive scaling. Large $v_t$ shrinks not just the data gradient but also the decay term — regularization effect becomes inconsistent across parameters.
AdamW (Loshchilov & Hutter): decouple weight decay from the gradient-based update.
In code, two lines in step():
if weight_decay != 0.0:
p.data -= lr * weight_decay * p.data
p.data -= lr * m_hat / (np.sqrt(v_hat) + eps)
Transformers almost universally use AdamW. GPT-2, LLaMA, my AL-1 pretrain — same pattern. Typical weight_decay=0.1 on weights; bias vectors and LayerNorm scales usually get weight_decay=0.
| Adam (L2 in grad) | AdamW | |
|---|---|---|
| Weight decay | Mixed into $g_t$ | Direct shrink on $\theta$ |
| Effect with large $v_t$ | Weaker regularization | Decay independent of $v_t$ |
| Default for LLMs | Legacy | Yes |
Default hyperparameters and what they mean
| Hyperparam | Typical value | Role |
|---|---|---|
lr |
$3 \times 10^{-4}$ | Global step scale (GPT-style pretrain) |
| $\beta_1$ | 0.9 | Gradient momentum window ~10 steps |
| $\beta_2$ | 0.999 | Squared-gradient window ~1000 steps |
| $\epsilon$ | $10^{-8}$ | Floor on $\sqrt{v}$ |
weight_decay |
0.1 | Decoupled L2 shrink per step |
Learning rate is the knob I still tune. Warmup + cosine decay come later in the AL-1 checklist; the optimizer itself doesn’t care — it just consumes whatever $\eta$ I pass each step.
$\beta_2 = 0.999$ means $v_t$ is a very slow average. Early training steps have tiny $\hat{v}_t$ after correction, so effective learning rates can be large until $v$ accumulates. That’s one reason warmup helps: small $\eta$ early while $v$ calibrates.
From-scratch AdamW (minimal)
Matches my AL-1 zyn/optim.py:
class AdamW:
def __init__(self, params, lr=3e-4, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.0):
self.params = list(params)
self.lr = lr
self.beta1, self.beta2 = betas
self.eps = eps
self.weight_decay = weight_decay
self.t = 0
self.m = [np.zeros_like(p.data) for p in self.params]
self.v = [np.zeros_like(p.data) for p in self.params]
def zero_grad(self):
for p in self.params:
p.grad = np.zeros_like(p.data)
def step(self):
self.t += 1
bc1 = 1.0 - self.beta1 ** self.t
bc2 = 1.0 - self.beta2 ** self.t
for i, p in enumerate(self.params):
g = p.grad
self.m[i] = self.beta1 * self.m[i] + (1.0 - self.beta1) * g
self.v[i] = self.beta2 * self.v[i] + (1.0 - self.beta2) * (g * g)
m_hat = self.m[i] / bc1
v_hat = self.v[i] / bc2
if self.weight_decay != 0.0:
p.data -= self.lr * self.weight_decay * p.data
p.data -= self.lr * m_hat / (np.sqrt(v_hat) + self.eps)
No gradcheck on the optimizer itself — I verify by overfitting a tiny model: one batch, loss should drop monotonically for 50 steps if CE + AdamW are wired correctly.
How does the full training step look?
sequenceDiagram
participant Data
participant Model
participant Loss
participant Autograd
participant AdamW
Data->>Model: batch (x, y)
Model->>Loss: logits
Loss->>Autograd: scalar L
Autograd->>Model: ∂L/∂θ
AdamW->>Model: θ ← θ − update
One step:
optimizer.zero_grad()- Forward → cross-entropy loss
loss.backward()- (Optional) clip global grad norm
optimizer.step()
I run gradcheck on the loss once. I sanity-check the optimizer by overfitting. Only then do I point it at a real corpus.
When AdamW is the wrong tool
Adam family excels at noisy, high-dimensional, non-stationary objectives — exactly language modeling. It can generalize slightly worse than tuned SGD on some vision tasks with careful schedules. For small convex problems SGD+momentum is simpler and sufficient.
Memory cost: two extra full-sized buffers per parameter ($m$ and $v$). For a 124M-param GPT that’s hundreds of MB of optimizer state on top of weights and activations. Tradeoff I accept.
Mistakes I watch for
| Bug | What happens |
|---|---|
Forgetting zero_grad() |
Gradients accumulate across steps |
| Weight decay on biases/LayerNorm | Shrinks scales that should stay near 1 |
| Adam L2 instead of AdamW | Regularization fights adaptive scaling |
| LR too high without warmup | Loss spikes, $v_t$ never recovers |
| Bias correction off-by-one ($t$ starts at 0) | First hundred steps subtly wrong |
Adam looked like a bag of heuristics the first time I read the paper. Momentum for direction, RMSprop for scale, bias correction for cold start, decay on the side — each piece solves a specific failure mode of SGD. AdamW is the version that plays nice with weight decay when training billion-parameter models.
I’ve got autograd, gradcheck, cross-entropy, and AdamW. The next items on my AL-1 checklist are gradient clipping, LR warmup, and a training loop that actually runs overnight without me babysitting the loss curve. The optimizer won’t save a wrong loss or a broken backward pass — but once those are solid, AdamW is what makes the update step boring in the best way.
Support the writing
If this post helped, a coffee keeps the deep dives, paper picks, and news digests coming.
Buy Me a Coffee