Arc 1: Mathematical & Computational Prerequisites

Optimizers: Momentum, Adam, and Learning Rate Schedules

12 min readbuilds on Gradients and How Machines Learn

SGDMomentumAdamx² + 10y²

Three update rules on the same bowl, with separately chosen learning rates: SGD 0.085, momentum 0.012, Adam 0.18. These paths illustrate the rules, not a performance ranking.

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$ . In stochastic gradient descent, we estimate that gradient from a sampled batch.

A single learning rate can make progress slow when the loss curves sharply in one direction and gently in another.

We will use a narrow valley to compare momentum, adaptive scaling, and AdamW, then put together a learning-rate schedule.

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.

Loss curvature can vary substantially across directions in parameter space. SGD multiplies every gradient coordinate by the same learning rate, so that rate must accommodate the steep directions.

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 slowly.

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:

f(x,y) = x² + 10y²

SGDL = 31.50

MomentumL = 31.50

AdamL = 31.50

The same bowl and starting point, with separately chosen learning rates: SGD 0.08, momentum 0.01, Adam 0.15. Compare update paths, not a controlled speed benchmark.

SGD alternates across the valley; momentum carries information from previous gradients; Adam also scales each coordinate using recent squared gradients.

Momentum

Vanilla SGD uses only the current gradient. Momentum accumulates a running velocity, so earlier gradients continue to influence the next update.

The velocity vector $v$ accumulates past gradients:

v_t = \beta \cdot v_{t-1} + \nabla L(\theta_t)

\theta_{t+1} = \theta_t - \eta \cdot v_t

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$ , gradients roughly ten steps back still contribute. This velocity is an unnormalized weighted sum; multiplying it by $1-\beta$ would give an exponential moving average.

In a direction where gradients alternate sign, their contributions to velocity can partly cancel. Where they consistently point the same way, they accumulate. The learning rate and momentum coefficient determine whether this improves progress or produces further overshoot.

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 computes the gradient at its own position (faint red) but updates use $-v$ (blue), which averages out the oscillation and points roughly along the floor:

SGD · lr=0.085

Momentum · lr=0.011, β=0.9

step 0 / 70

Each ball evaluates the gradient at its own location. The blue arrow includes the accumulated velocity; the red arrow uses the current gradient alone.

Polyak introduced momentum in 1964; Chapter 8 of Deep Learning develops its connection to a damped physical system.

Per-parameter adaptive rates: RMSProp

Momentum still uses a single learning rate. RMSProp additionally rescales each coordinate using its recent gradient magnitudes.

Hinton presented RMSProp in his 2012 Coursera lectures. It tracks a running average of squared gradients:

s_t = \beta \cdot s_{t-1} + (1 - \beta) \cdot (\nabla L)^2

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

\theta_{t+1} = \theta_t - \frac{\eta}{\sqrt{s_t} + \epsilon} \cdot \nabla L(\theta_t)

Parameters with consistently large gradients get their effective learning rate shrunk. Parameters with consistently small gradients get their effective learning rate boosted. The $\epsilon$ (usually $10^{-8}$ ) prevents division by zero and also affects updates when the denominator is small.

$s_t$ has units of squared gradient. Taking $\sqrt{s_t}$ puts the denominator on the gradient scale, so $g_t / \sqrt{s_t}$ compares the current gradient with recent magnitudes.

Adam: combining the two moments

Adam, introduced by Kingma and Ba in 2014, combines first- and second-moment estimates.

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)

m_t = \beta_1 \cdot m_{t-1} + (1 - \beta_1) \cdot g_t

v_t = \beta_2 \cdot v_{t-1} + (1 - \beta_2) \cdot g_t^2

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:

\hat{m}_t = \frac{m_t}{1 - \beta_1^t}, \quad \hat{v}_t = \frac{v_t}{1 - \beta_2^t}

Then the update is:

\theta_{t+1} = \theta_t - \eta \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}

Common defaults are $\beta_1 = 0.9$ , $\beta_2 = 0.999$ , and $\epsilon = 10^{-8}$ ; their suitability depends on the run.

This example compares two proportional gradient histories: A's values are around 5 and B's are A's divided by 100. In Adam's update $\eta \cdot \hat{m} / (\sqrt{\hat{v}} + \epsilon)$ , this scale factor cancels between the numerator and denominator apart from epsilon. The resulting step sizes are similar despite the different raw gradients.

Inspect update 1 of 12

Parameter A

g5.200e+0

m̂5.200e+0

√v̂5.200e+0

|Δθ|1.000e-2

Parameter B

g5.200e-2

m̂5.200e-2

√v̂5.200e-2

|Δθ|1.000e-2

Shared scales: g, m̂, and √v̂ run from 0 to 6; |Δθ| runs from 0 to 0.01.

Both columns use the same scales. Step through the proportional histories to compare their updates.

A worked Adam step

For three consecutive steps on a single parameter, let $\theta$ start at 1.0 and use gradients 0.5, 0.4, and 0.6. We use η = 0.1, β₁ = 0.9, β₂ = 0.999, and ε = 10⁻⁸:

η = 0.1 · β₁ = 0.9 · β₂ = 0.999 · ε = 10⁻⁸

Gradient g: 0.5000; chosen input gradient
First moment m: 0.0500; 0.9 m(previous) + 0.1 g
Second moment v: 2.500e-4; 0.999 v(previous) + 0.001 g²
Corrected m̂: 0.5000; m / (1 − 0.9^1)
Corrected v̂: 0.2500; v / (1 − 0.999^1)
Update Δθ: -0.1000; −0.1 m̂ / (√v̂ + 1e−8)
New parameter θ: 0.9000; 1.0000 + (-0.1000)

Each selected step includes the moments accumulated during earlier steps. At step 1, correction takes m from 0.05 to 0.5 and v from 0.00025 to 0.25.

Select a step to inspect both moment updates, their bias corrections, and the resulting parameter. Each step includes the history of the earlier ones.

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 update about $31.6\eta$ , too large relative to $\eta$ ; correcting only $v$ would make it $0.1\eta$ , too small, because $\sqrt{v_1} \approx 0.0158$ is much smaller than $\sqrt{\hat{v}_1} = 0.5$ in the denominator.

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.

AdamW: decoupled weight decay

When training neural networks, you almost always want some form of regularization to prevent overfitting. The most common is weight decay, where you gently shrink selected weights toward zero (biases and normalization parameters are often excluded). 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.

In Adam, adding the L2 term to the gradient also changes the first- and second-moment estimates. That produces a different update from shrinking the parameters directly. AdamW separates these operations so the decay term does not pass through the moment estimates.

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

\theta_{t+1} = \theta_t - \eta \left( \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda \theta_t \right)

The weight decay term $\lambda \theta_t$ shrinks parameters directly, outside Adam's moment scaling. An optimizer described as "Adam" may implement this decoupled form or apply L2 regularization through the gradient; inspect the update rule when reproducing a training recipe.

Learning rate schedules: warmup + cosine decay

A learning-rate schedule changes the size of updates during training. Warmup limits the initial steps; decay makes later updates smaller.

Warmup

Large initial updates can destabilize training before the optimizer has useful gradient statistics. Warmup reduces these early updates while those statistics and the model activations change.

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

\eta_t = \eta_{\text{max}} \cdot \frac{t}{T_{\text{warmup}}} \quad \text{for } t \leq T_{\text{warmup}}

The warmup length is a hyperparameter, not a fixed percentage. The original LLaMA paper used 2,000 warmup steps. The right choice depends on the model, batch size, and peak learning rate.

Cosine decay

A decay phase gradually reduces the learning rate toward a chosen minimum. One common choice is cosine decay:

\eta_t = \eta_{\text{min}} + \frac{1}{2}(\eta_{\text{max}} - \eta_{\text{min}}) \left(1 + \cos\left(\frac{t - T_{\text{warmup}}}{T_{\text{total}} - T_{\text{warmup}}} \cdot \pi\right)\right)

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). It decays slowly at first, accelerates in the middle, and slows down near the end.

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):

phase

warmup

step

0 / 30000

learning rate

3.00e-6

An illustrative 30k-step schedule with 3k warmup steps. The dot marks the current step.

The cosine schedule appeared in work on warm restarts and was later used in Chinchilla and LLaMA training. It requires a chosen decay horizon. Warmup-Stable-Decay schedules instead keep a stable phase and begin a cooldown when training is nearing its end.

What this costs at scale

Adam and AdamW store two moment tensors, $m$ and $v$ , with one entry per parameter. With fp32 moments, that is 8 bytes per parameter, in addition to weights, gradients, and activations.

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.

With bf16 weights (2 bytes per parameter), two fp32 moment tensors occupy four times as many bytes as the weights. An optional fp32 master-weight copy adds another 4 bytes per parameter. The exact allocation depends on the training implementation. ZeRO can shard these states across devices, and low-bit optimizers reduce their storage precision.

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)$ :

Rosenbrock: (1 − x)² + 100(y − x²)²

SGDL = 104.00

AdamL = 104.00

Rosenbrock trajectories with SGD at 0.0005 and Adam at 0.01. This illustration caps gradient norm at 40; the paths depend on those choices.

The practical recipe

Here is a complete PyTorch example of AdamW with warmup followed by cosine decay:

import torch
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR
 
torch.manual_seed(0)
model = torch.nn.Linear(1, 1)
x = torch.linspace(-1, 1, 32).unsqueeze(1)
y = 3 * x + 1
steps, warmup_steps = 100, 10
optimizer = AdamW(model.parameters(), lr=0.05, weight_decay=0.01)
warmup = LinearLR(optimizer, start_factor=0.01, total_iters=warmup_steps)
cosine = CosineAnnealingLR(optimizer, T_max=steps - warmup_steps, eta_min=0.005)
scheduler = SequentialLR(optimizer, [warmup, cosine], milestones=[warmup_steps])
 
for step in range(steps):
    optimizer.zero_grad(set_to_none=True)
    loss = torch.nn.functional.mse_loss(model(x), y)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()
    scheduler.step()
print(f"Final loss: {loss.item():.4f}")

The code fits a line so you can run the whole loop. Its learning rate and ten-step warmup are for this toy problem. A transformer recipe needs its own tuning; the reusable part is the order of backward, clipping, optimizer step, and scheduler step.

Common misconceptions

"Adam always beats SGD." Results depend on the objective, tuning, data, and compute budget. AdamW is a common transformer optimizer, but that convention is not a guarantee of better convergence or generalization.

"Learning rate is the only hyperparameter that matters." Warmup, weight decay, clipping, and the beta values affect the updates too. For example, a lower β₂ makes the second-moment estimate respond more quickly to changing gradients. Their useful values depend on the training run.

"One recipe works for every transformer." AdamW with warmup and cosine decay is a useful recipe to understand. Peak learning rate, batch size, decay, clipping, betas, and the schedule still need evaluation for the actual run. Other optimizers and schedules make different memory and convergence tradeoffs.

What's next

We have a training loop and several ways to turn a gradient into a parameter update. The next question is how to store and compute those updates when memory and arithmetic precision are limited.

The next post is GPUs, Floating Point, and Why Precision Matters: what fp32, bf16, and fp8 represent, and why training combines formats.

Additional reading (and watching)

Kingma, D. P. & Ba, J. (2014). Adam: A Method for Stochastic Optimization. ICLR 2015.
Loshchilov, I. & Hutter, F. (2017). Decoupled Weight Decay Regularization. ICLR 2019.
Loshchilov, I. & Hutter, F. (2016). SGDR: Stochastic Gradient Descent with Warm Restarts. ICLR 2017.
Goodfellow, I., Bengio, Y. & Courville, A. (2016). Deep Learning, Ch. 8: Optimization for Training Deep Models. MIT Press.
Ruder, S. (2016). An Overview of Gradient Descent Optimization Algorithms.
Karpathy, A. (2019). A Recipe for Training Neural Networks.
Hoffmann, J., et al. (2022). Training Compute-Optimal Large Language Models (Chinchilla).
Dettmers, T., et al. (2021). 8-bit Optimizers via Block-wise Quantization.