Optimizers: Momentum, Adam, and Learning Rate Schedules

Optimizers like Momentum, Adam, and learning rate schedules solve problems with vanilla SGD, such as zigzagging in steep valleys and slow progress in flat directions. Momentum adds inertia to smooth oscillations, while Adam adapts per-parameter learning rates for faster convergence. These techniques are essential for training large models like 70B-parameter LLMs.

Optimizers: Momentum, Adam, and Learning Rate Schedules Gradients and How Machines Learn /blog/gradients-and-learning Three optimizers, same start, same valley. SGD bounces off the steep walls, momentum glides, Adam adapts per-parameter and arrives first. In the last post /blog/gradients-and-learning , we got gradient descent working. We computed gradients, applied the chain rule backwards through a network, and updated parameters with a simple rule: . 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 , an elongated bowl where the -direction is ten times steeper than the -direction. Watch SGD bounce: Same starting point, three optimizers. SGD oscillates, momentum smooths the path, Adam adapts per-parameter and converges fastest. Watch all three run. SGD red oscillates wildly in the -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 that accumulates past gradients: The hyperparameter typically 0.9 controls how much of the previous velocity we retain. When , this is just vanilla SGD. When , 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 , which flips wildly across the valley walls. On the right, momentum keeps the same gradient faint red, still flipping but updates use blue , which averages out the oscillation and points roughly along the floor: Both balls feel the same gradient at every step. The momentum ball ignores the wobble and follows the running average instead. 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 : 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 usually is just there to prevent division by zero. The square root matters here. has units of squared gradient , so puts the denominator back on the same scale as the gradient itself, and the ratio 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 , the effective step sizes for A and B end up at the same order of magnitude: Same η, same betas, same machinery. But because each parameter gets divided by its own √v̂, the effective step size lands in the same ballpark for both. That's per-parameter adaptation in one picture. 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 : a running average of the gradients like momentum's velocity Second moment : a running average of the squared gradients like RMSProp's denominator There's one wrinkle. Both and 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 , , 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. starts at 1.0. The gradients arriving from the loss are 0.5, then 0.4, then 0.6. Standard Adam hyperparameters throughout: Three Adam updates on one parameter. The yellow row is the quantity being computed in this phase. The footer note tracks how aggressively bias correction is rescuing the estimate from its zero-init. The thing to notice on step 1: the raw , but . The second moment is corrected the same way: , but . Both corrections work together; they're not redundant. Correcting only would make the first step too small relative to its target scale; correcting only would make it too large, because is much smaller than 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 -correction with becomes negligible within a hundred steps or so. The -correction with takes much longer: at it's still about , and even at it's about . 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 to the gradient is the same as adding 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 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 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 down to often set to something like 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 decay, blue : Warmup ramps from near-zero to η max over the first ~10% of training, then cosine decay cools the rate down to η min. Real LLM pretraining usually uses a much shorter warmup often well under 1% of total steps ; the demo exaggerates the warmup phase so the ramp is visually obvious. 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 and the second moment . 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 :