Optimizers in deep learning: SGD, momentum, Adam and AdamW

Overview

An optimizer is the part of training that changes a model's weights. After backpropagation has computed the gradient, which says how the loss would change if each weight moved a little, the optimizer decides how far and in which direction to actually move each weight. It does this again on every batch, often hundreds of thousands of times.

The simplest rule is to step a little way downhill along the gradient. Practical optimizers add memory and adaptivity. Momentum keeps a running direction so the model does not zigzag. Adaptive methods such as Adam give each weight its own effective step size based on how large and noisy its gradients have been. Weight decay gently pulls weights towards zero to discourage overfitting. A learning-rate schedule changes the overall step size over the course of training.

An everyday comparison is walking down a foggy hillside to reach the valley. You can only feel the slope under your feet. Plain gradient descent takes a step straight downhill each time. Momentum is like a heavy ball that keeps rolling through small bumps. Adam is like a hiker who takes shorter steps on steep, slippery ground and longer ones on gentle, steady slopes. The learning-rate schedule is the decision to stride out at first and take careful steps near the bottom.

What it is

Most optimizers in use today are built from a few ingredients.

Optimizer or ingredient	What it does	Typical use
SGD	Moves weights by the learning rate times the gradient	Baseline; vision models with momentum
Momentum	Keeps a moving average of past gradients as the step direction	Smooths noisy updates, speeds progress
Nesterov momentum	Looks ahead along the momentum direction before computing the step	Slightly better-behaved momentum
AdaGrad	Scales steps by the accumulated history of squared gradients	Sparse features; shrinks steps over time
RMSprop	Scales steps by a moving average of squared gradients	Recurrent networks, reinforcement learning
Adam	Combines momentum with RMSprop-style scaling and bias correction	General default for many models
AdamW	Adam with weight decay applied separately from the gradient step	Standard for transformers and large models
Learning-rate schedule	Changes the learning rate over training, such as warm-up then cosine decay	Used with almost every optimizer
Gradient clipping	Caps the size of the gradient before the update	Stabilises transformers and RNNs

Optimizers have state. SGD with momentum stores one extra number per weight; Adam and AdamW store two. For a model with a billion parameters, Adam's state alone adds two billion numbers, which is why memory-saving techniques such as DeepSpeed ZeRO, FSDP and GaLore focus heavily on optimizer state.

Where it came from

Stochastic approximation, the idea behind stochastic gradient descent, goes back to Herbert Robbins and Sutton Monro in 1951. Boris Polyak described the momentum method in 1964, and Yurii Nesterov published his accelerated method in 1983. These ideas were used to train neural networks from the 1980s onward.

Adaptive methods arrived in the 2010s. AdaGrad was published by John Duchi, Elad Hazan and Yoram Singer in 2011. RMSprop was proposed by Geoffrey Hinton in his online lectures in 2012. Diederik Kingma and Jimmy Ba introduced Adam in a 2014 preprint, published at ICLR 2015, with default settings that worked well across many problems.

Ilya Loshchilov and Frank Hutter then showed that the usual way of adding L2 regularisation to Adam did not behave like true weight decay, and proposed AdamW, which decouples the two. Their paper was published at ICLR 2019, and AdamW became the standard choice for transformers. The same authors had earlier popularised cosine learning-rate schedules with warm restarts. Research continues, and newer optimizers such as Lion and Muon have been proposed and tried on large models, but AdamW remains the common default.

How it works

Plain SGD

For each weight w with gradient g and learning rate η, SGD updates w to w − η × g. "Stochastic" means the gradient comes from a small random batch rather than the full dataset, which is noisy but cheap.

Momentum

Momentum keeps a velocity v. Each step, v becomes β × v + g, with β often 0.9, and the weight moves by −η × v. Consistent gradient directions build up speed; directions that flip back and forth cancel out.

Adam

Adam keeps two moving averages per weight: m, the average gradient, and v, the average squared gradient. Typical decay rates are β1 = 0.9 and β2 = 0.999. Because both start at zero, Adam applies a bias correction in early steps. The update is the corrected m divided by the square root of the corrected v plus a small epsilon, times the learning rate. Weights with large, noisy gradients take smaller steps; weights with small, consistent gradients take relatively larger ones.

AdamW

AdamW performs the Adam step and then separately shrinks each weight by a small amount: w becomes w − η × λ × w, where λ is the weight decay coefficient, often around 0.01 to 0.1. Keeping this separate from the adaptive scaling makes weight decay behave as intended.

Learning-rate schedules

A common recipe for transformers is linear warm-up for the first few hundred or thousand steps, then cosine decay towards a small final value. Warm-up avoids large, unstable updates while Adam's averages are still forming. Decay lets the model settle into a good region at the end.

In PyTorch, a typical setup looks like this:

import torch
opt = torch.optim.AdamW(model.parameters(), lr=3e-4, betas=(0.9, 0.95), weight_decay=0.1)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=10_000)

Key concepts and terms

Term	Meaning
Learning rate	Overall step size for weight updates
Batch size	Number of examples used to compute each gradient
Momentum	Moving average of gradients used as the update direction
Adaptive learning rate	Per-weight step sizes based on gradient history
Bias correction	Adam's adjustment for averages that start at zero
Weight decay	Shrinking weights towards zero each step
Decoupled weight decay	Applying weight decay separately from the gradient step, as in AdamW
Warm-up	Gradually raising the learning rate at the start of training
Cosine decay	Lowering the learning rate along a cosine curve
Gradient clipping	Limiting the gradient's size before updating
Optimizer state	Extra per-weight values an optimizer stores, such as Adam's averages
Convergence	Reaching a point where further training brings little improvement

A worked example

This example uses simple numbers to compare SGD and Adam on one weight. The loss is L = w², so the gradient is 2w and the best value is w = 0. Start at w = 3.

SGD with learning rate 0.1:

Step 1: gradient = 6, so w becomes 3 − 0.6 = 2.4.
Step 2: gradient = 4.8, so w becomes 2.4 − 0.48 = 1.92.
Each step multiplies w by 0.8. After 10 steps w is about 0.32.

Adam with learning rate 0.1:

Step 1: m and v after bias correction equal 6 and 36. The update is 0.1 × 6 / √36 = 0.1, so w becomes 2.9.
Early on, Adam moves each weight by roughly the learning rate per step, whatever the gradient's size.
After 10 steps w is near 2.0, a slower start here, but on real problems with gradients of very different sizes across weights, this size-independence is exactly what helps.

Now an illustrative training scenario. A team fine-tunes a 300-million-parameter language model on 2 million tokens per epoch. They choose AdamW with a peak learning rate of 0.0002, 200 warm-up steps, cosine decay over 5,000 steps, weight decay of 0.1 and gradient clipping at a norm of 1.0. When loss spikes at step 900, they lower the peak learning rate by half and the spikes disappear.

Strengths and limitations

Strengths:

SGD with momentum is simple, memory-light and generalises well for many vision models.
Adam and AdamW work well across architectures with little tuning.
Decoupled weight decay gives predictable regularisation.
Schedules and warm-up make training stable at large scale.

Limitations:

No optimizer is best everywhere; results depend on the model and data.
Adam-family optimizers need extra memory for state.
The learning rate remains sensitive and usually needs tuning.
Changing batch size often requires retuning the learning rate.
Newer optimizers may show gains in some papers that do not carry over to every setting.

How it compares with alternatives

Optimizer	Memory per weight	Strength	Watch out for
SGD with momentum	One extra value	Strong generalisation on vision	Needs careful learning-rate tuning
Adam	Two extra values	Fast, robust default	Weight decay behaves poorly if added as L2
AdamW	Two extra values	Standard for transformers	Memory cost at very large scale
RMSprop	One extra value	Works well for recurrent and RL models	Less common today
Bayesian optimisation of hyperparameters	Not an optimizer of weights	Tunes learning rate and decay	Can be paired with any optimizer

Practical guidance and best practices

Start with AdamW for transformers and most new models, and SGD with momentum for convolutional image models trained from scratch.
Tune the learning rate first; it matters more than most other choices.
Use warm-up followed by a decay schedule, such as cosine, for large models.
Exclude biases and normalisation layer parameters from weight decay, as many recipes do.
Clip gradients by global norm, often at 1.0, for transformers.
When you raise the batch size, raise the learning rate too, then re-check stability.
Log the learning rate and gradient norm alongside the loss so problems are visible.
Use tools such as Optuna or Ray Tune for systematic tuning when stakes are high.

Common mistakes

Using Adam with an L2 penalty and expecting it to act like AdamW's weight decay.
Keeping one learning rate when switching from SGD to Adam; their good ranges differ widely.
Skipping warm-up on large models and getting early loss spikes or divergence.
Forgetting to step the scheduler, or stepping it per epoch when it was designed per step.
Restarting training from a checkpoint without restoring optimizer state, which causes a jolt in loss.

Frequently asked questions

Which optimizer should I use?

For most transformer and general deep learning work, AdamW with warm-up and cosine decay is a strong default. For convolutional image models trained from scratch, SGD with momentum remains a good choice and often generalises well.

What is the difference between Adam and AdamW?

Both use the same adaptive update. Adam implementations often add weight decay as a term in the gradient, which then gets scaled by the adaptive factors. AdamW applies weight decay directly to the weights, separately from the gradient step, which works better in practice.

Why do I need a learning-rate schedule?

Early in training, large steps make fast progress, though the very first steps benefit from warm-up. Later, smaller steps help the model settle. A schedule automates this and usually improves final results.

What learning rate should I start with?

Common starting points are around 0.001 for Adam on small models, and between 0.0001 and 0.0006 for AdamW on transformers. For SGD with momentum on image models, values near 0.1 are common. Always confirm with a few short trial runs.

Why does Adam use so much memory?

It stores two moving averages for every parameter, so its state is twice the size of the model itself. Sharding the state across GPUs or using memory-efficient variants reduces this.

Glossary

Epoch: one full pass through the training data.
Gradient: the direction and rate of steepest increase of the loss.
Hyperparameter: a setting such as the learning rate that is chosen, not learned.
L2 regularisation: adding the sum of squared weights to the loss.
Loss landscape: the shape of the loss as a function of all weights.
Minibatch: a small random subset of data used for one update.
Step: one update of the weights.
Velocity: the running direction kept by momentum methods.

Summary

Optimizers turn gradients into weight updates. Plain SGD steps downhill; momentum smooths the path; Adam adapts each weight's step size; and AdamW adds properly decoupled weight decay, which made it the standard for transformers. Combined with warm-up, decay schedules and gradient clipping, these methods train nearly every modern network, while the learning rate remains the most important setting to tune.

It belongs to Building blocks in my stack, under Deep learning. Others in the same group:

Attention mechanism Backpropagation Normalization layers Dropout Word embeddings U-Net Knowledge distillation Pruning