Intermediate
15 min read
#Optimizers#Momentum#RMSProp#Adam#AdamW#AdaGrad#Adaptive Learning

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:

  1. 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.
  2. 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.
  3. 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.
🐍 Python
import 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 Formulation
v_t = beta * v_(t-1) + gradient_t
theta_new = theta_old - learning_rate * v_t

Where:

  • v_t is the "velocity" (an exponentially weighted moving average of past gradients)
  • beta is the momentum coefficient (commonly set to 0.9), controlling how much of the previous velocity is retained
  • v_0 is 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.

🐍 Python
import 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):

🐍 Python
import 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 Formulation
lookahead_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.

🐍 Python
import 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):

🐍 Python
import 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 Formulation
G_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_t is the running sum of squared gradients for that specific parameter, accumulated across all time steps so far
  • epsilon is 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.

🐍 Python
import 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):

🐍 Python
import 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 Formulation
E[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]_t is the exponentially decaying moving average of squared gradients (typically decay_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.

🐍 Python
import 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):

🐍 Python
import 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 Formulation
m_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.

🐍 Python
import 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):

🐍 Python
import 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 Formulation
m_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.

🐍 Python
import 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):

🐍 Python
import 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#

OptimizerKey Idea AddedSolvesMain Limitation
SGDBasic gradient-based updateEstablishes the baselineSame LR for all parameters, no memory of past gradients, prone to zig-zagging
MomentumRunning average of gradients ("velocity")Zig-zagging, slow convergence in ravinesCan overshoot the minimum
Nesterov MomentumLook-ahead gradient before applying momentumMomentum's overshoot problemStill uses a single global learning rate
AdaGradPer-parameter adaptive learning rate (accumulated squared gradients)Single global LR being suboptimal for sparse/uneven dataLearning rate shrinks forever and can stall training
RMSPropDecaying average of squared gradients instead of accumulating sumAdaGrad's ever-shrinking learning rateNo use of gradient direction memory (no momentum)
AdamCombines Momentum (1st moment) + RMSProp (2nd moment) + bias correctionLack of both direction-smoothing and adaptive scaling togetherWeight decay coupling distorts regularization
AdamWDecouples weight decay from the adaptive gradient updateAdam's distorted weight decay/regularization behaviorStill has multiple hyperparameters to tune (betas, eps, weight_decay)

Quick Recap (Beginner to Advanced Flow)#

  1. 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.
  2. Momentum adds a running average ("velocity") of past gradients to smooth the path and accelerate convergence in consistent directions.
  3. Nesterov Momentum improves on Momentum by computing the gradient at a "look-ahead" position, correcting the trajectory earlier and reducing overshoot.
  4. 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.
  5. RMSProp fixes AdaGrad's ever-shrinking learning rate by using a decaying moving average of squared gradients instead of an ever-growing sum.
  6. 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.
  7. 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.
Knowledge Checkpoint

06. Optimizers Checkpoint

Q1.How does Polyak Momentum accelerate gradient descent in ravines and noisy directions?
AIt accumulates an exponentially decaying moving average of past velocity vectors v_t = gamma*v_{t-1} + lr*g_t, dampening oscillations and accelerating along consistent gradients.
BIt doubles the learning rate after every successful batch.
CIt replaces matrix multiplications with dot products.
DIt disables backpropagation on alternating steps.
Q2.What distinguishes the Adam optimizer from RMSProp and AdaGrad?
AAdam combines both first-moment estimation (momentum/mean) and second-moment estimation (uncentered variance/RMSProp) with initialization bias-correction factors.
BAdam computes true Hessian second-order matrix inversions.
CAdam does not require learning rates.
DAdam only operates on recurrent networks.
Q3.Why is AdamW preferred over standard Adam when training modern Transformer and deep architectures with weight decay?
AAdamW decouples L2 weight decay from the gradient update, directly shrinking weights theta = theta - lr * lambda * theta rather than letting moving variance distort regularization.
BAdamW runs 5x faster on CPUs.
CAdamW removes the need for layer normalization.
DAdamW sets all gradients to 1.
Track Your Learning

Finished studying this notebook?

Mark this guide as completed to update your course progress roadmap.