06. Advanced Optimizers (SGD, Momentum, RMSProp, Adam, AdamW)
In-depth guide to modern deep learning optimizers: SGD with Momentum, Nesterov, AdaGrad, RMSProp, Adam, and decoupled weight decay AdamW.
Optimizers: Complete Notes (Beginner to Advanced)
Introduction#
An optimizer is the algorithm responsible for updating a neural network's weights and biases using the gradients computed by backpropagation. Plain gradient descent (covered in the previous topic) is the simplest optimizer, but it has practical limitations: it can converge slowly, get stuck oscillating in narrow valleys of the loss surface, and use the same learning rate for every parameter regardless of how frequently or rarely that parameter's gradient changes.
The optimizers below were developed, one building on the ideas of the previous, to address these specific limitations. This note walks through them roughly in the order they were historically introduced, so each one's motivation makes sense in context.
1. SGD#
Stochastic Gradient Descent, in the context of optimizers, refers to the basic update rule applied per mini-batch:
theta_new = theta_old - learning_rate * gradient
As covered in the Gradient Descent notes, "SGD" is used loosely in frameworks to mean this basic update rule applied to mini-batches (not strictly one sample at a time). It is the foundation on which every other optimizer in this note builds.
Limitations of plain SGD that motivate the rest of this note:
- Same learning rate for every parameter: every weight in the network is updated using the exact same step size, even though some parameters may need larger updates and others may need much smaller, more careful updates.
- No memory of past gradients: each update only looks at the current gradient, ignoring the direction the optimizer has been moving in previous steps. This makes it prone to zig-zagging in narrow, curved loss surfaces (also called "ravines"), and slow to build up speed in consistently flat, gently-sloped directions.
- Sensitive to noisy gradients: since mini-batch gradients are noisy estimates of the true gradient, plain SGD can be pushed around by this noise, especially with small batch sizes.
🐍 PythonInteractive WebAssemblyimport torch
import torch.nn as nn
model = nn.Linear(10, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
# Each optimizer.step() call applies: theta = theta - lr * gradient
2. Momentum#
Momentum addresses SGD's zig-zagging and slow convergence problem by giving the optimizer a form of "memory" of previous update directions, similar to a ball rolling down a hill that builds up speed and doesn't immediately change direction just because the slope changes slightly.
Update rule:
Mathematical Formulationv_t = beta * v_(t-1) + gradient_t theta_new = theta_old - learning_rate * v_t
Where:
v_tis the "velocity" (an exponentially weighted moving average of past gradients)betais the momentum coefficient (commonly set to 0.9), controlling how much of the previous velocity is retainedv_0is typically initialized to 0
Intuition: instead of updating based purely on the current gradient, momentum accumulates a running average of past gradients. If the gradient consistently points in a similar direction across several steps, the velocity builds up, and the optimizer moves faster and faster in that consistent direction. If the gradient direction fluctuates (as in a zig-zagging ravine), the fluctuating components partially cancel out in the running average, smoothing the overall path.
Without momentum: in loss surfaces shaped like narrow ravines (steep in one direction, shallow in another), plain SGD tends to oscillate back and forth across the steep direction while making very slow progress along the shallow direction toward the actual minimum.
With momentum: oscillations across the steep direction get dampened (since they cancel out in the running average across time), while progress along the consistently-sloped shallow direction accelerates (since the velocity keeps building up in that direction), leading to noticeably faster and smoother convergence.
🐍 PythonInteractive WebAssemblyimport torch
import torch.nn as nn
model = nn.Linear(10, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
# Internally applies: v = 0.9*v_prev + gradient; theta = theta - lr*v
Code example (momentum from scratch):
🐍 PythonInteractive WebAssemblyimport numpy as np
def f(x):
return x[0]**2 + 10 * x[1]**2 # a "ravine" shaped loss: steep in y, shallow in x
def grad_f(x):
return np.array([2*x[0], 20*x[1]])
def gradient_descent_with_momentum(start, lr=0.05, beta=0.9, steps=30):
x = np.array(start, dtype=float)
v = np.zeros_like(x)
path = [x.copy()]
for _ in range(steps):
grad = grad_f(x)
v = beta * v + grad
x = x - lr * v
path.append(x.copy())
return path
path = gradient_descent_with_momentum(start=[5.0, 5.0])
print("Final position:", path[-1])
3. Nesterov Momentum#
Nesterov Momentum (also called Nesterov Accelerated Gradient, or NAG) is a refinement of standard momentum. Instead of computing the gradient at the current position and then applying momentum, it computes the gradient at the position the momentum is about to move to (a "look-ahead" position), giving it a chance to correct course slightly before fully committing to the step.
Update rule:
Mathematical Formulationlookahead_position = theta_old - learning_rate * beta * v_(t-1) v_t = beta * v_(t-1) + gradient_at(lookahead_position) theta_new = theta_old - learning_rate * v_t
Intuition: standard momentum first computes the gradient exactly where it currently stands, then takes a big momentum-boosted step. This means momentum can sometimes overshoot the minimum before it "notices" (via the gradient) that it needs to slow down or change direction. Nesterov momentum instead peeks ahead to where the momentum term is about to carry it, computes the gradient there, and uses that look-ahead gradient to adjust the step. This gives the optimizer a chance to "correct" its trajectory earlier, since it's effectively evaluating the slope at its predicted future position rather than only its current position.
Without Nesterov momentum: standard momentum can overshoot minima more often, since it commits to the momentum-driven step before checking what the gradient looks like at the destination, leading to more oscillation before eventually settling.
With Nesterov momentum: the optimizer is more responsive to changes in the loss surface's curvature ahead of its current path, generally leading to faster convergence and less overshoot compared to standard momentum, particularly for convex optimization problems.
🐍 PythonInteractive WebAssemblyimport torch
import torch.nn as nn
model = nn.Linear(10, 1)
# Setting nesterov=True enables Nesterov momentum instead of standard momentum
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9, nesterov=True)
Code example (Nesterov momentum from scratch):
🐍 PythonInteractive WebAssemblyimport numpy as np
def grad_f(x):
return np.array([2*x[0], 20*x[1]])
def nesterov_momentum(start, lr=0.05, beta=0.9, steps=30):
x = np.array(start, dtype=float)
v = np.zeros_like(x)
path = [x.copy()]
for _ in range(steps):
lookahead = x - lr * beta * v # peek ahead using current velocity
grad_lookahead = grad_f(lookahead) # compute gradient AT the lookahead position
v = beta * v + grad_lookahead
x = x - lr * v
path.append(x.copy())
return path
path = nesterov_momentum(start=[5.0, 5.0])
print("Final position:", path[-1])
4. AdaGrad#
AdaGrad (Adaptive Gradient Algorithm) introduces a fundamentally different idea: instead of using the same learning rate for every parameter, it adapts the learning rate individually for each parameter, based on how large that parameter's past gradients have been.
Update rule:
Mathematical FormulationG_t = G_(t-1) + gradient_t^2 (element-wise, accumulated per parameter) theta_new = theta_old - (learning_rate / (sqrt(G_t) + epsilon)) * gradient_t
Where:
G_tis the running sum of squared gradients for that specific parameter, accumulated across all time steps so farepsilonis a tiny constant (e.g.,1e-8) added purely to prevent division by zero- The division happens element-wise, meaning each parameter gets its own effective learning rate based on its own gradient history
Intuition: parameters that have received large gradients frequently (meaning G_t grows large) get their effective learning rate shrunk more aggressively, since the loss surface is "steep and sensitive" in that direction, needing smaller, more careful steps. Parameters that have received small or infrequent gradients (meaning G_t stays small) keep a relatively larger effective learning rate, allowing them to keep making meaningful progress.
Where this is especially useful: AdaGrad performs well on sparse data (e.g., text data with word embeddings, where some parameters/features are updated very rarely while others are updated constantly), since it naturally boosts the learning rate for those rarely-updated parameters relative to frequently-updated ones.
Without AdaGrad's per-parameter adaptation: using a single global learning rate for every parameter forces a compromise. A rate suited for frequently-updated, high-gradient parameters might be far too large for rarely-updated, low-gradient parameters (and vice versa), leading to inefficient training on data where different features/parameters are updated at very different frequencies.
The major limitation of AdaGrad (which motivates RMSProp below): because G_t is a continuously accumulating sum of squared gradients that only ever grows larger over time, the effective learning rate learning_rate / sqrt(G_t) continuously shrinks and can eventually become so tiny that learning effectively stops, even if the model hasn't yet reached a good solution. This is a significant practical problem for training over many epochs.
🐍 PythonInteractive WebAssemblyimport torch
import torch.nn as nn
model = nn.Linear(10, 1)
optimizer = torch.optim.Adagrad(model.parameters(), lr=0.01)
Code example (AdaGrad from scratch):
🐍 PythonInteractive WebAssemblyimport numpy as np
def grad_f(x):
return np.array([2*x[0], 20*x[1]])
def adagrad(start, lr=0.5, epsilon=1e-8, steps=30):
x = np.array(start, dtype=float)
G = np.zeros_like(x) # accumulated squared gradients, per parameter
path = [x.copy()]
for _ in range(steps):
grad = grad_f(x)
G += grad ** 2
x = x - (lr / (np.sqrt(G) + epsilon)) * grad
path.append(x.copy())
return path
path = adagrad(start=[5.0, 5.0])
print("Final position:", path[-1])
5. RMSProp#
RMSProp (Root Mean Square Propagation) fixes AdaGrad's main weakness (the ever-shrinking, never-recovering learning rate) by using an exponentially decaying moving average of squared gradients instead of a simple, ever-growing running sum.
Update rule:
Mathematical FormulationE[g^2]_t = decay_rate * E[g^2]_(t-1) + (1 - decay_rate) * gradient_t^2 theta_new = theta_old - (learning_rate / (sqrt(E[g^2]_t) + epsilon)) * gradient_t
Where:
E[g^2]_tis the exponentially decaying moving average of squared gradients (typicallydecay_rate = 0.9)- Because this is a moving average (not an ever-accumulating sum), older gradients gradually "fade out" of the average rather than permanently inflating it forever
Intuition: RMSProp still adapts the learning rate per parameter based on the magnitude of recent gradients, just like AdaGrad, but because it only remembers a decaying window of recent gradient history rather than the entire training history, the effective learning rate can go back up again if recent gradients become smaller, instead of monotonically shrinking toward zero forever.
Without this decaying average (i.e., sticking with AdaGrad's approach): training over long periods (many epochs) risks the effective learning rate becoming so small that learning stalls out prematurely, even if the loss landscape ahead still has room for meaningful improvement.
With RMSProp's decaying average: the per-parameter learning rate stays responsive to the current local behavior of the gradients rather than the entire accumulated training history, allowing the optimizer to keep adapting effectively throughout long training runs without prematurely grinding to a halt.
🐍 PythonInteractive WebAssemblyimport torch
import torch.nn as nn
model = nn.Linear(10, 1)
optimizer = torch.optim.RMSprop(model.parameters(), lr=0.01, alpha=0.9)
# 'alpha' here is the decay_rate in the formula above
Code example (RMSProp from scratch):
🐍 PythonInteractive WebAssemblyimport numpy as np
def grad_f(x):
return np.array([2*x[0], 20*x[1]])
def rmsprop(start, lr=0.1, decay_rate=0.9, epsilon=1e-8, steps=30):
x = np.array(start, dtype=float)
E_g2 = np.zeros_like(x) # decaying average of squared gradients
path = [x.copy()]
for _ in range(steps):
grad = grad_f(x)
E_g2 = decay_rate * E_g2 + (1 - decay_rate) * (grad ** 2)
x = x - (lr / (np.sqrt(E_g2) + epsilon)) * grad
path.append(x.copy())
return path
path = rmsprop(start=[5.0, 5.0])
print("Final position:", path[-1])
6. Adam#
Adam (Adaptive Moment Estimation) is currently the most widely used optimizer in deep learning. It combines the two key ideas covered so far: Momentum (tracking a moving average of the gradient itself, called the "first moment") and RMSProp (tracking a moving average of the squared gradient, called the "second moment"), then uses both together to compute the update.
Update rule:
Mathematical Formulationm_t = beta1 * m_(t-1) + (1 - beta1) * gradient_t (first moment: mean of gradients, like Momentum) v_t = beta2 * v_(t-1) + (1 - beta2) * gradient_t^2 (second moment: mean of squared gradients, like RMSProp) m_hat_t = m_t / (1 - beta1^t) (bias-corrected first moment) v_hat_t = v_t / (1 - beta2^t) (bias-corrected second moment) theta_new = theta_old - learning_rate * m_hat_t / (sqrt(v_hat_t) + epsilon)
Where common default values are beta1 = 0.9, beta2 = 0.999, and epsilon = 1e-8. The variable t refers to the current time step (starting from 1), used specifically for the bias correction step.
Why bias correction is needed: m_t and v_t are both initialized to 0. In the early steps of training, this initialization at zero biases the moving averages toward zero (they haven't had enough steps yet to reflect the true average of the gradients seen so far). The bias correction terms (1 - beta1^t and 1 - beta2^t) counteract this early-step bias; as t grows large, beta1^t and beta2^t approach 0, so the correction factor approaches 1 and has little effect later in training, but it matters significantly in the first several steps.
Intuition (putting the pieces together):
- The first moment (
m_t) acts like Momentum, smoothing out the direction of the update using a running average of the raw gradient, helping accelerate convergence and dampen oscillation. - The second moment (
v_t) acts like RMSProp, adaptively scaling the learning rate per parameter based on the recent magnitude of that parameter's gradients, giving smaller effective steps to parameters with consistently large gradients and larger effective steps to parameters with consistently small gradients.
Without combining these two ideas (using only Momentum or only RMSProp separately): you would get the acceleration and oscillation-dampening benefits of Momentum without adaptive per-parameter learning rates, or the adaptive per-parameter learning rates of RMSProp without the smoothing/acceleration benefit of tracking the gradient's own direction over time. You would have to pick one benefit or the other.
With Adam: you get both benefits simultaneously, which is a major reason Adam tends to work well as a strong default optimizer across a very wide range of problems with relatively little manual tuning, making it extremely popular in practice.
🐍 PythonInteractive WebAssemblyimport torch
import torch.nn as nn
model = nn.Linear(10, 1)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-8)
Code example (Adam from scratch):
🐍 PythonInteractive WebAssemblyimport numpy as np
def grad_f(x):
return np.array([2*x[0], 20*x[1]])
def adam(start, lr=0.1, beta1=0.9, beta2=0.999, epsilon=1e-8, steps=30):
x = np.array(start, dtype=float)
m = np.zeros_like(x)
v = np.zeros_like(x)
path = [x.copy()]
for t in range(1, steps + 1):
grad = grad_f(x)
m = beta1 * m + (1 - beta1) * grad
v = beta2 * v + (1 - beta2) * (grad ** 2)
m_hat = m / (1 - beta1 ** t)
v_hat = v / (1 - beta2 ** t)
x = x - lr * m_hat / (np.sqrt(v_hat) + epsilon)
path.append(x.copy())
return path
path = adam(start=[5.0, 5.0])
print("Final position:", path[-1])
7. AdamW#
AdamW is a modification of Adam that changes how weight decay (L2-style regularization) is applied. The name stands for "Adam with decoupled Weight decay."
The problem AdamW fixes: in the original Adam formulation, if you want to add L2 regularization, the standard approach is to add the regularization term directly into the gradient before it goes through Adam's momentum and adaptive scaling calculations:
gradient_with_L2 = gradient + weight_decay * theta
This gradient (now including the weight decay term) is then used to compute m_t and v_t exactly as in standard Adam. The problem is that this couples the weight decay effect with Adam's adaptive per-parameter learning rate scaling (division by sqrt(v_hat_t)), which distorts the weight decay's intended effect: parameters with large historical gradients (large v_t) end up having their weight decay effectively shrunk by that same adaptive scaling, when weight decay is supposed to uniformly pull all weights toward zero regardless of their gradient history.
AdamW's fix: it "decouples" weight decay from the gradient-based moment calculations entirely. Instead of folding weight decay into the gradient, it applies the weight decay directly to the parameter during the final update step, separately from the Adam update:
Mathematical Formulationm_t = beta1 * m_(t-1) + (1 - beta1) * gradient_t v_t = beta2 * v_(t-1) + (1 - beta2) * gradient_t^2 m_hat_t = m_t / (1 - beta1^t) v_hat_t = v_t / (1 - beta2^t) theta_new = theta_old - learning_rate * ( m_hat_t / (sqrt(v_hat_t) + epsilon) + weight_decay * theta_old )
Notice the weight_decay * theta_old term is added outside and independently of the m_hat_t / (sqrt(v_hat_t) + epsilon) term, rather than being mixed into the gradient before those moment calculations happen. This is the "decoupling."
Without decoupled weight decay (standard Adam + L2 regularization mixed into the gradient): the actual amount of regularization effectively applied to each parameter ends up being inconsistent and distorted by that same parameter's adaptive learning rate scaling, weakening the regularization's intended, uniform effect and making it harder to tune weight decay independently of the learning rate.
With AdamW's decoupled weight decay: the regularization strength behaves consistently and predictably as intended (uniformly shrinking weights toward zero), independent of each parameter's individual gradient history, making weight decay easier to tune correctly and generally leading to better generalization performance. This is why AdamW has largely replaced plain Adam as the default optimizer for training many modern large-scale models, especially transformers.
🐍 PythonInteractive WebAssemblyimport torch
import torch.nn as nn
model = nn.Linear(10, 1)
# weight_decay here is decoupled, applied as in AdamW's formula above
optimizer = torch.optim.AdamW(model.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01)
Code example (AdamW from scratch):
🐍 PythonInteractive WebAssemblyimport numpy as np
def grad_f(x):
return np.array([2*x[0], 20*x[1]])
def adamw(start, lr=0.1, beta1=0.9, beta2=0.999, epsilon=1e-8, weight_decay=0.01, steps=30):
x = np.array(start, dtype=float)
m = np.zeros_like(x)
v = np.zeros_like(x)
path = [x.copy()]
for t in range(1, steps + 1):
grad = grad_f(x) # NOTE: weight decay is NOT added into this gradient
m = beta1 * m + (1 - beta1) * grad
v = beta2 * v + (1 - beta2) * (grad ** 2)
m_hat = m / (1 - beta1 ** t)
v_hat = v / (1 - beta2 ** t)
# Weight decay applied directly and separately to the parameter itself
x = x - lr * (m_hat / (np.sqrt(v_hat) + epsilon) + weight_decay * x)
path.append(x.copy())
return path
path = adamw(start=[5.0, 5.0])
print("Final position:", path[-1])
Summary Table: Optimizer Evolution#
| Optimizer | Key Idea Added | Solves | Main Limitation |
|---|---|---|---|
| SGD | Basic gradient-based update | Establishes the baseline | Same LR for all parameters, no memory of past gradients, prone to zig-zagging |
| Momentum | Running average of gradients ("velocity") | Zig-zagging, slow convergence in ravines | Can overshoot the minimum |
| Nesterov Momentum | Look-ahead gradient before applying momentum | Momentum's overshoot problem | Still uses a single global learning rate |
| AdaGrad | Per-parameter adaptive learning rate (accumulated squared gradients) | Single global LR being suboptimal for sparse/uneven data | Learning rate shrinks forever and can stall training |
| RMSProp | Decaying average of squared gradients instead of accumulating sum | AdaGrad's ever-shrinking learning rate | No use of gradient direction memory (no momentum) |
| Adam | Combines Momentum (1st moment) + RMSProp (2nd moment) + bias correction | Lack of both direction-smoothing and adaptive scaling together | Weight decay coupling distorts regularization |
| AdamW | Decouples weight decay from the adaptive gradient update | Adam's distorted weight decay/regularization behavior | Still has multiple hyperparameters to tune (betas, eps, weight_decay) |
Quick Recap (Beginner to Advanced Flow)#
- SGD is the basic update rule: subtract the learning rate times the gradient. It uses a single global learning rate and no memory of past gradients.
- Momentum adds a running average ("velocity") of past gradients to smooth the path and accelerate convergence in consistent directions.
- Nesterov Momentum improves on Momentum by computing the gradient at a "look-ahead" position, correcting the trajectory earlier and reducing overshoot.
- AdaGrad introduces per-parameter adaptive learning rates based on the accumulated sum of squared gradients, but this sum only grows, causing the learning rate to shrink toward zero over time.
- RMSProp fixes AdaGrad's ever-shrinking learning rate by using a decaying moving average of squared gradients instead of an ever-growing sum.
- Adam combines Momentum's gradient-averaging (first moment) with RMSProp's squared-gradient-averaging (second moment), plus bias correction, making it a strong, widely-used default optimizer.
- AdamW refines Adam by decoupling weight decay from the adaptive gradient calculations, applying it directly to the parameters instead, leading to more effective and predictable regularization, and has become the standard choice for training many modern large-scale models.
06. Optimizers Checkpoint
Finished studying this notebook?
Mark this guide as completed to update your course progress roadmap.