What is the difference between SGD, momentum, and Adam?

Plain SGD steps in the direction of the negative gradient scaled by one learning rate shared across every parameter: theta -= lr * grad. Momentum adds memory — it keeps an exponentially decaying running sum of past gradients (the "velocity") and steps along that, so consistent directions accelerate and noisy ones cancel, like a ball rolling downhill. Adam adds a second kind of memory: alongside the momentum-like first moment (mean of the gradients) it tracks the second moment (mean of the squared gradients) per parameter, and divides the step by the square root of that. The effect is a separate, self-tuning learning rate for every weight — parameters with large, noisy gradients get small steps; quiet parameters get large ones. On an ill-conditioned loss surface SGD is forced to use a tiny learning rate to stay stable in the steep directions, which starves the shallow ones; Adam rescales each direction so it can make even progress everywhere.

Why do transformers need learning-rate warmup?

Because Adam's per-parameter learning rate is only trustworthy once it has seen enough gradients to estimate the second moment. In the first handful of steps that estimate is built from almost no samples, so — as the RAdam paper showed — the adaptive learning rate "has problematically large variance in the early stage." A full-size step computed from a wildly uncertain denominator can throw the weights into a bad region the model never recovers from. Warmup ramps the learning rate up from near zero over the first few hundred to few thousand steps, which "works as a variance reduction technique": it keeps the early, unreliable steps small until the moment estimates settle. This is also exactly why resuming training with fresh (zeroed) optimizer state causes a loss spike — you have thrown away the warmed-up estimates.

What is the difference between L2 regularization and weight decay (AdamW)?

For plain SGD they are the same thing rescaled by the learning rate, but for Adam they are not — and that is the whole point of AdamW. Classic L2 adds lambda*theta to the gradient before Adam's adaptive rescaling, so the decay on each weight gets divided by the square root of that weight's second moment: weights with large gradients get decayed less, which is backwards. AdamW decouples the decay — it applies lambda*theta directly to the weights, outside the adaptive step: theta -= lr * (mhat/(sqrt(vhat)+eps) + lambda*theta). Every weight then decays at the same rate regardless of its gradient history. The paper shows this "decouples the optimal choice of weight decay factor from the setting of the learning rate" and "substantially improves Adam's generalization." AdamW is the default for transformers for exactly this reason.

Cosine, linear, or WSD learning-rate schedule — which should I use?

All three warm up, then decay; they differ in the decay shape and in one practical property. Cosine decay (eta_t = eta_min + 0.5*(eta_max-eta_min)*(1+cos(pi*t/T))) spends a long time near the peak and decays smoothly to the floor; it is the long-standing default and hard to beat when you know the total step count T in advance. Linear decay is simpler and, empirically, nearly as good. WSD (Warmup-Stable-Decay) holds the learning rate constant at its peak for most of training and only decays in a short final phase — its advantage is that you do not have to commit to T up front: you can keep training at the stable rate and trigger the decay whenever you decide to stop, and it matches cosine's final loss. What matters most is not which curve you pick but that you warm up, that the peak LR is right, and that you actually decay to a small fraction of it by the end.

What happens if you resume training from a checkpoint without the optimizer state?

The loss spikes upward and then slowly recovers — it usually does not crash. When you restore the weights but reset Adam's first and second moments to zero, the next few steps are computed from moment estimates built on almost no data, and bias correction makes those first steps effectively large and mis-scaled — the same instability warmup was invented to avoid, now happening mid-training. The optimizer re-estimates the moments over the following tens to hundreds of steps and the loss comes back down, but you have wasted compute and perturbed the trajectory. The fix is to always checkpoint and restore the full optimizer state (m and v for every parameter) plus the learning-rate schedule position, not just the weights. This is why an Adam checkpoint is roughly three times the size of the raw model.

Optimization Dynamics: Why Adam, Why Warmup, Why Cosine

An optimizer never sees the loss surface — it gets one gradient at a time and has to turn that stream of local hints into a path to the bottom. We race SGD, momentum, and Adam down the same dialable ill-conditioned valley (real update rules, trajectories drawn live), then build Adam precisely — per-parameter learning rates from the second moment, bias correction honestly — and AdamW's decoupled decay. Then the practitioner's layer: why warmup exists, what cosine/linear/WSD actually buy, gradient clipping and loss spikes, muP transfer, and batch-size-vs-LR — all computed, all animated.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of Optimization Dynamics: Why Adam, Why Warmup, Why Cosine — built to be understood, not skimmed.

Questions

What is the difference between SGD, momentum, and Adam?: Plain SGD steps in the direction of the negative gradient scaled by one learning rate shared across every parameter: theta -= lr * grad. Momentum adds memory — it keeps an exponentially decaying running sum of past gradients (the "velocity") and steps along that, so consistent directions accelerate and noisy ones cancel, like a ball rolling downhill. Adam adds a second kind of memory: alongside the momentum-like first moment (mean of the gradients) it tracks the second moment (mean of the squared gradients) per parameter, and divides the step by the square root of that. The effect is a separate, self-tuning learning rate for every weight — parameters with large, noisy gradients get small steps; quiet parameters get large ones. On an ill-conditioned loss surface SGD is forced to use a tiny learning rate to stay stable in the steep directions, which starves the shallow ones; Adam rescales each direction so it can make even progress everywhere.
Why do transformers need learning-rate warmup?: Because Adam's per-parameter learning rate is only trustworthy once it has seen enough gradients to estimate the second moment. In the first handful of steps that estimate is built from almost no samples, so — as the RAdam paper showed — the adaptive learning rate "has problematically large variance in the early stage." A full-size step computed from a wildly uncertain denominator can throw the weights into a bad region the model never recovers from. Warmup ramps the learning rate up from near zero over the first few hundred to few thousand steps, which "works as a variance reduction technique": it keeps the early, unreliable steps small until the moment estimates settle. This is also exactly why resuming training with fresh (zeroed) optimizer state causes a loss spike — you have thrown away the warmed-up estimates.
What is the difference between L2 regularization and weight decay (AdamW)?: For plain SGD they are the same thing rescaled by the learning rate, but for Adam they are not — and that is the whole point of AdamW. Classic L2 adds lambda*theta to the gradient before Adam's adaptive rescaling, so the decay on each weight gets divided by the square root of that weight's second moment: weights with large gradients get decayed less, which is backwards. AdamW decouples the decay — it applies lambda*theta directly to the weights, outside the adaptive step: theta -= lr * (mhat/(sqrt(vhat)+eps) + lambda*theta). Every weight then decays at the same rate regardless of its gradient history. The paper shows this "decouples the optimal choice of weight decay factor from the setting of the learning rate" and "substantially improves Adam's generalization." AdamW is the default for transformers for exactly this reason.
Cosine, linear, or WSD learning-rate schedule — which should I use?: All three warm up, then decay; they differ in the decay shape and in one practical property. Cosine decay (eta_t = eta_min + 0.5*(eta_max-eta_min)*(1+cos(pi*t/T))) spends a long time near the peak and decays smoothly to the floor; it is the long-standing default and hard to beat when you know the total step count T in advance. Linear decay is simpler and, empirically, nearly as good. WSD (Warmup-Stable-Decay) holds the learning rate constant at its peak for most of training and only decays in a short final phase — its advantage is that you do not have to commit to T up front: you can keep training at the stable rate and trigger the decay whenever you decide to stop, and it matches cosine's final loss. What matters most is not which curve you pick but that you warm up, that the peak LR is right, and that you actually decay to a small fraction of it by the end.
What happens if you resume training from a checkpoint without the optimizer state?: The loss spikes upward and then slowly recovers — it usually does not crash. When you restore the weights but reset Adam's first and second moments to zero, the next few steps are computed from moment estimates built on almost no data, and bias correction makes those first steps effectively large and mis-scaled — the same instability warmup was invented to avoid, now happening mid-training. The optimizer re-estimates the moments over the following tens to hundreds of steps and the loss comes back down, but you have wasted compute and perturbed the trajectory. The fix is to always checkpoint and restore the full optimizer state (m and v for every parameter) plus the learning-rate schedule position, not just the weights. This is why an Adam checkpoint is roughly three times the size of the raw model.

Optimization Dynamics: Why Adam, Why Warmup, Why Cosine

Questions

Related explainers