Optimizers: Momentum, Adam, and Learning Rate Schedules

In the last post, we got gradient descent working. We computed gradients, applied the chain rule backwards through a network, and updated parameters with a simple rule: $\theta \leftarrow \theta - \eta \nabla L$. Vanilla stochastic gradient descent. The most straightforward optimizer you can write.

And it works, in theory. On toy problems, SGD finds the minimum just fine. The moment you try to train something real with it, you run into problems.

This post walks through those problems and the sequence of clever fixes that got us to the optimizer almost everyone uses today to train large language models: AdamW with warmup and cosine decay. The path from "SGD works in theory" to "here's what actually trains a 70B parameter model" is surprisingly short, and each step along the way is motivated by a concrete failure of the previous approach. (If you want a wider survey of every optimizer variant ever proposed, Ruder's overview is the standard reference.)

The problem with vanilla SGD

Imagine you're walking down a long, skinny valley. The walls are steep, but the floor slopes gently toward the minimum at the far end. SGD bounces back and forth between the steep walls while barely making progress along the floor. Each step overshoots in the steep direction and undershoots in the shallow one. You end up zigzagging when you should be gliding.

This is not contrived. Loss surfaces in deep learning are full of dimensions like this. Some parameters have enormous gradients (the loss is very sensitive to them), others have tiny gradients (the loss barely cares). SGD treats them all the same. One learning rate, one step size, take it or leave it.

If you set the learning rate high enough to make progress in the flat direction, you'll overshoot and oscillate in the steep direction. Set it low enough to be stable in the steep direction, and you'll crawl along the valley floor at geological speeds.

Here's what that looks like. The surface below is $f(x,y) = x^2 + 10y^2$, an elongated bowl where the $y$-direction is ten times steeper than the $x$-direction. Watch SGD bounce:

Watch all three run. SGD (red) oscillates wildly in the $y$-direction. Momentum (blue) dampens those oscillations and builds up speed along the valley. Adam (green) does both and gets there first. Same starting point, same surface, dramatically different paths.

How do we fix this? Three ideas, layered on top of each other.

Fix 1: momentum

The first insight is almost physical. If SGD is like a ball rolling downhill, vanilla SGD is a ball with no mass. It instantly changes direction with every gradient. There's no inertia. The fix is to give the ball some weight.

Instead of using just the current gradient to update the parameters, we maintain a velocity vector $v$ that accumulates past gradients:

The hyperparameter $\beta$ (typically 0.9) controls how much of the previous velocity we retain. When $\beta = 0$, this is just vanilla SGD. When $\beta = 0.9$, we're essentially averaging over the last ~10 gradients.

Think about what this does to our zigzagging problem. In the steep direction, the gradient flips sign every step (positive, negative, positive, negative). When you average those out, the oscillations cancel. In the flat direction, the gradient consistently points the same way, step after step. Those accumulate. The velocity builds up, and you start moving faster and faster along the valley floor.

Side by side, here are the two arrows that matter at every step. On the left, vanilla SGD: the descent direction is just $-\nabla L$, which flips wildly across the valley walls. On the right, momentum keeps the same gradient (faint red, still flipping) but updates use $-v$ (blue), which averages out the oscillation and points roughly along the floor:

The gradient oscillations in noisy dimensions get dampened. The consistent signal in important dimensions gets amplified. Momentum was introduced by Polyak back in 1964, and nearly sixty years later it's still a core ingredient in every serious optimizer. (Goodfellow et al.'s Deep Learning Chapter 8 has a great derivation if you want the full mechanical picture.)

Fix 2: per-parameter adaptive rates (RMSProp)

Momentum helps with oscillation, but it doesn't solve the fundamental issue that different parameters live in different worlds. Some weights in a neural network change the loss enormously with a tiny perturbation. Others barely register. Using the same learning rate for all of them is like paying every employee the same salary regardless of their job.

The idea behind RMSProp (Hinton, 2012, introduced in a Coursera lecture with no paper, which is great) is to track how big each parameter's gradients have been historically, and then scale the learning rate inversely.

Maintain a running average of squared gradients:

Then divide the update by $\sqrt{s_t}$:

Parameters with consistently large gradients get their effective learning rate shrunk. Parameters with consistently small gradients get their effective learning rate boosted. The optimizer is adapting on a per-parameter basis. The $\epsilon$ (usually $10^{-8}$) is just there to prevent division by zero.

The square root matters here. $s_t$ has units of squared gradient, so $\sqrt{s_t}$ puts the denominator back on the same scale as the gradient itself, and the ratio $g_t / \sqrt{s_t}$ is dimensionally a unitless "this gradient relative to recent gradient size." Drop the square root and the units don't line up; the denominator overcompensates and the update gets vanishingly small for any parameter with even moderate gradient magnitude.

This matters a lot in practice. Instead of the human having to pick one learning rate that somehow works for every parameter in a model with billions of weights, the optimizer figures out the right scale for each one automatically.

Here's that adaptation made literal. Two parameters, side by side. Parameter A's gradient is consistently around 5; parameter B's is around 0.05, a hundred times smaller. Watch what Adam does to each. The raw gradients are wildly different, but after the running averages settle and we compute $\eta \cdot \hat{m} / (\sqrt{\hat{v}} + \epsilon)$, the effective step sizes for A and B end up at the same order of magnitude:

Fix 3: Adam (combining both)

Adam, introduced by Kingma and Ba in 2014, does the obvious thing: combine momentum and RMSProp.

It maintains two running averages:

• First moment $m_t$: a running average of the gradients (like momentum's velocity)
• Second moment $v_t$: a running average of the squared gradients (like RMSProp's denominator)

There's one wrinkle. Both $m$ and $v$ are initialized at zero. In the early steps, they're biased toward zero because they haven't accumulated enough history. Adam corrects for this with a bias correction step:

Then the update is:

The default hyperparameters ($\beta_1 = 0.9$, $\beta_2 = 0.999$, $\epsilon = 10^{-8}$) are surprisingly robust. People often just use them as-is.

So Adam gives you momentum's oscillation dampening plus RMSProp's per-parameter adaptation, with a bias correction to keep things honest in early training. It's a clean design, and it's been the default optimizer for most of deep learning since about 2015.

A worked Adam step

The bias correction is one of those things that looks like algebraic noise until you watch it happen on real numbers. So let's walk through three consecutive steps on a single parameter. $\theta$ starts at 1.0. The gradients arriving from the loss are 0.5, then 0.4, then 0.6. Standard Adam hyperparameters throughout:

The thing to notice on step 1: the raw $m_1 = 0.1 \cdot 0.5 = 0.05$, but $\hat{m}_1 = 0.05 / (1 - 0.9) = 0.5$. The second moment is corrected the same way: $v_1 = 0.001 \cdot 0.25 = 0.00025$, but $\hat{v}_1 = 0.00025 / (1 - 0.999) = 0.25$. Both corrections work together; they're not redundant. Correcting only $m$ would make the first step too small relative to its target scale; correcting only $v$ would make it too large, because $\sqrt{v_1} \approx 0.0158$ is much smaller than $\sqrt{\hat{v}_1} = 0.5$ in the denominator. With both corrections in place, the first Adam update lands at the intended scale.

The two correction factors fade at different rates. The $m$-correction $1/(1-\beta_1^t)$ with $\beta_1 = 0.9$ becomes negligible within a hundred steps or so. The $v$-correction $1/(1-\beta_2^t)$ with $\beta_2 = 0.999$ takes much longer: at $t=100$ it's still about $1/0.095 \approx 10.5$, and even at $t=1000$ it's about $1/0.632 \approx 1.58$. Bias correction matters for thousands of steps for the second moment, not just the first few.

But there's a subtle issue with how Adam handles regularization.

AdamW: weight decay done right

When training neural networks, you almost always want some form of regularization to prevent overfitting. The most common is weight decay, where you gently push all weights toward zero. In vanilla SGD, weight decay and L2 regularization are mathematically equivalent: adding $\lambda \theta$ to the gradient is the same as adding $\frac{\lambda}{2} \|\theta\|^2$ to the loss.

But in Adam, they're not equivalent. Because Adam rescales the gradient by the second moment, the weight decay term gets rescaled too. Large weights that happen to have large gradients get less decay than they should, and small weights with small gradients get more. The regularization is tangled up with the adaptive learning rate, and neither works as intended.

Loshchilov and Hutter (2017) pointed this out and proposed AdamW, which decouples the weight decay from the gradient-based update:

The weight decay term $\lambda \theta_t$ is applied directly to the parameters, not passed through Adam's adaptive scaling machinery. It's a small change in the code, but it matters a lot empirically. In modern transformer pretraining, AdamW with decoupled weight decay is effectively the default, and "Adam" in casual speech often means AdamW in practice. Papers and code aren't always consistent about it, though, and other parts of the field (vision fine-tuning, Adafactor for memory-constrained T5-style training, newer optimizers like Lion or Shampoo) really do use something else. When it matters, check the implementation.

Learning rate schedules: warmup + cosine decay

Picking the right learning rate is important, but here's the thing: the right learning rate changes over the course of training. Early on, the model's weights are randomly initialized and the loss landscape is chaotic. Late in training, the model has found a good basin and needs to fine-tune. Using the same learning rate the whole time is suboptimal.

Warmup

At initialization, the model is completely random. The gradients are noisy and potentially huge. If you slam the optimizer with a large learning rate from step one, you can blow up the training entirely. The loss spikes, the gradients explode, and you're dead.

The fix is simple: start with a very small learning rate and linearly increase it to your target over the first $N$ warmup steps.

For LLM pretraining, warmup is typically 1 to 2% of total training steps. LLaMA used 2,000 warmup steps for example, which is a fraction of a percent of its total training. It's not doing anything fancy. It's just giving the model time to stabilize before you hit the gas. Karpathy calls this "letting the network sort out its initial chaos," and that's basically the right way to think about it.

Cosine decay

After warmup, you want to gradually reduce the learning rate toward zero. The most popular schedule for LLM pretraining is cosine decay:

The formula looks busy but the idea is simple. It traces a smooth half-cosine from $\eta_{\text{max}}$ down to $\eta_{\text{min}}$ (often set to something like $0.1 \times \eta_{\text{max}}$ or even zero). The key feature is that it decays slowly at first, accelerates in the middle, and slows down near the end. It gives the optimizer plenty of time at a reasonable learning rate during the bulk of training, then gradually cools things off.

Here is the whole schedule end to end. The dot rides through 30k steps: a fast linear ramp during the first 10% of training (warmup, amber), then a smooth cosine descent down toward $\eta_\text{min}$ (decay, blue):

Why cosine specifically? Honestly, the original motivation was about warm restarts, but in practice cosine decay just works well. It became a strong default for LLM pretraining largely because the Chinchilla and LLaMA papers used it. Schedule design isn't fully settled, though. Work on Warmup-Stable-Decay (WSD) and related cooldown schedules is an active area, partly because cosine requires committing to a fixed training horizon up front while WSD keeps a long stable phase and only decays near the end.

The full picture is to ramp up linearly during warmup, then cosine decay down. If you plot the learning rate over training, it looks like a sharp rise followed by a smooth descent.

What this costs at scale

Here's a practical detail that matters a lot at scale. Adam (and AdamW) maintains two extra tensors for every parameter in the model: the first moment $m$ and the second moment $v$. So the optimizer state is 2x the size of the model itself.

For a 7B parameter model stored in fp32 (4 bytes per parameter), the model weights take 28 GB. The optimizer state takes another 56 GB. For a 70B model? 280 GB of weights plus 560 GB of optimizer state. This is why training large models requires so much more memory than inference, even before you account for gradients and activations.

In real mixed-precision training the exact multiplier shifts around. Most stacks keep parameters in bf16 but maintain fp32 master weights and fp32 Adam moments alongside, which still puts optimizer state at roughly 2× the model size in bytes. Distributed-training schemes like ZeRO shard the moments across devices so any single GPU only holds a fraction. The big picture is unchanged: Adam state is one of the dominant memory costs, and that's what makes techniques like 8-bit optimizers (Dettmers et al. 2021), which quantize the moments, so important for fitting training on limited hardware.

A harder landscape

The elongated bowl above is a clean illustration, but real loss surfaces are gnarlier. Here's the Rosenbrock function, a classic optimization benchmark with a narrow, curved valley where the minimum sits at $(1, 1)$: