Advanced Optimizers: SGD, Momentum, RMSProp, Adam

Objectives

After this lesson, you will be able to:

• Identify the three fundamental weaknesses of standard gradient descent: slow progress through ill-conditioned ravines, oscillations driven by mini-batch noise, and a single shared learning rate across heterogeneous parameters.
• Formulate the update equations for Momentum, AdaGrad, RMSProp, and Adam, explaining the role of each constituent term.
• Explain why Adam requires bias correction, demonstrating its impact with step-1 calculations.
• Choose appropriate optimizers, base learning rates, and learning rate schedules when approaching new machine learning tasks.

Prerequisites

• Gradient Descent, in Machine Learning (Learned)Machine Learning

Quick Reminder

Gradient Descent demonstrated that on quadratic surfaces, GD converges if $\eta < 2/L$, where $L$ is maximum curvature. When directional curvatures differ sharply—yielding a large condition number $\kappa$—contours form elongated ellipses, causing standard GD to crawl. That lesson also introduced mini-batching, where gradients are estimated on subsets of size $B$ rather than the full dataset of size $m$.

Intuition

Imagine rolling a heavy ball down a steep, narrow ravine whose floor gently slopes toward one end. The ball does not roll straight down the valley floor. It oscillates back and forth against the steep valley walls several times before gradually progressing along the bottom. Unlike gradient descent, however, a physical ball possesses momentum: the longer it rolls in a consistent direction, the faster it travels.

Standard gradient descent lacks momentum entirely. At each step, it only evaluates the instantaneous local slope under its feet. Consider a synthetic toy ravine to observe this limitation:

$$\mathcal{L}(w_1, w_2) = \frac{1}{2}\left(w_1^2 + 100\,w_2^2\right)$$

The minimum sits at $(0, 0)$. Curvature is 1 along $w_1$ and 100 along $w_2$, yielding condition number $\kappa = 100$: the valley is 10 times longer than it is wide ($\sqrt{\kappa} = 10$). We initialize optimization at $(-8, 1.5)$.

Three Weaknesses of Standard Gradient Descent

1. Narrow ravines. The steep slope along $w_2$ forces a tight step bound $\eta < 2/100 = 0.02$. Setting $\eta = 0.018$ near the stability boundary, distance to the valley floor along $w_2$ scales by $1 - 0.018 \times 100 = -0.8$ each step: $w_2$ oscillates wildly back and forth across the ravine floor (zigzagging). Along $w_1$, distance scales by $1 - 0.018 \times 1 = 0.982$: each iteration advances by only 1.8%. Standard GD requires 369 iterations to enter a 0.01 radius around the origin, spending nearly all its steps crawling along the valley floor.

2. Mini-batch noise. In stochastic mini-batch settings, gradient estimates computed on small subsets fluctuate, causing search paths to jitter randomly. With fixed $\eta$, SGD wanders perpetually around the minimum rather than settling at the bottom.

3. A single global learning rate for all parameters. At the starting point, the gradient vector is $[-8, 150]^\top$: the gradient along $w_2$ is almost 19 times larger than along $w_1$. Under a single shared $\eta$, parameters with steep gradients take massive strides while flat parameters crawl, despite the flat parameter needing to cover far greater distance (from $-8$ to 0). In deep networks with millions of parameters, gradient magnitudes span several orders of magnitude.

Two Remedial Ideas

1. Retain directional history. Accumulate historical gradients into a persistent “velocity”. Along directions where gradients point consistently, velocity accumulates (accelerating down the ravine). Along directions where gradients oscillate signs, opposing components cancel out (damping transversal oscillation). This is Momentum.
2. Individualize learning rates. Scale each parameter’s step inversely by its typical historical gradient magnitude. Steep parameters are damped; flat parameters are boosted. This is the core mechanism of AdaGrad and RMSProp.

Adam unifies both concepts. The objective of this lesson is: engineer gradient update rules that navigate narrow ravines rapidly, dampen stochastic noise, and eliminate manual per-parameter learning rate tuning.

Theory

Unified Notation

All algorithms in this lesson build directly on gradient descent. We index iterations starting from $t = 1$:

$$\mathbf{g}_t = \nabla_{\boldsymbol{\theta}}\mathcal{L}\left(\boldsymbol{\theta}_{t-1}\right), \qquad t = 1, 2, 3, \dots$$

where $\mathbf{g}_t$ is the gradient computed at parameters $\boldsymbol{\theta}_{t-1}$. Standard GD is expressed as $\boldsymbol{\theta}_t = \boldsymbol{\theta}_{t-1} - \eta\,\mathbf{g}_t$. The optimizers below maintain internal auxiliary vectors of matching dimension, known as the optimizer state, initialized to $\mathbf{0}$.

From RMSProp onward, vector squares, square roots, and divisions operate element-wise: $\mathbf{g}_t^2$ denotes the vector of squared elements $g_j^2$, rather than a scalar dot product.

Momentum: Adding Velocity

Momentum (Polyak, 1964), also known as the heavy-ball method, maintains a persistent velocity vector $\mathbf{v}$:

$$\begin{aligned} \mathbf{v}_t &= \beta\,\mathbf{v}_{t-1} + \mathbf{g}_t \\ \boldsymbol{\theta}_t &= \boldsymbol{\theta}_{t-1} - \eta\,\mathbf{v}_t \end{aligned}$$

The coefficient $\beta \in [0, 1)$ is the momentum factor, typically set to $\beta = 0.9$. The first equation retains fraction $\beta$ of previous velocity and adds the current gradient. The second equation updates parameters along velocity rather than raw gradient. When $\beta = 0$, $\mathbf{v}_t = \mathbf{g}_t$, recovering standard GD.

What does velocity remember? Expanding $\mathbf{v}_{t-1}$ recursively down to $\mathbf{v}_0 = \mathbf{0}$:

$$\mathbf{v}_t = \mathbf{g}_t + \beta\,\mathbf{g}_{t-1} + \beta^2\,\mathbf{g}_{t-2} + \dots + \beta^{t-1}\,\mathbf{g}_1$$

Velocity is an exponentially decayed sum of all past gradients. Consider two illustrative regimes:

• Consistently aligned gradients (along the valley floor), $\mathbf{g}_t = \mathbf{g}$: the sum forms a geometric series $\mathbf{g}\left(1 + \beta + \beta^2 + \dots\right)$, converging to $\mathbf{g} / (1 - \beta)$. For $\beta = 0.9$, steps become 10 times longer than standard GD.
• Alternating gradients (across the ravine walls), $\mathbf{g}, -\mathbf{g}, \mathbf{g}, \dots$: the sum evaluates to $\mathbf{g}\left(1 - \beta + \beta^2 - \dots\right)$, converging to $\mathbf{g} / (1 + \beta) \approx 0.53\,\mathbf{g}$. Opposing components cancel out by roughly half.

On our toy ravine with $\eta = 0.018$, Momentum enters the 0.01 error ball in 111 steps versus 369 for GD. The trade-off is inertial overshoot: the heavy ball surges down the valley, overshooting the minimum to $w_1 \approx 2.19$ before reversing.

Alternative literature formulates velocity as an exponential moving average (EMA): $\mathbf{v}_t = \beta\,\mathbf{v}_{t-1} + (1 - \beta)\,\mathbf{g}_t$. This scales velocity down by $(1 - \beta)$, meaning both formulations produce identical trajectories if the latter scales $\eta$ up by $1 / (1 - \beta)$. Always verify which convention a codebase uses before adopting learning rates.

Nesterov accelerated gradient computes the gradient at the anticipated lookahead position $\boldsymbol{\theta}_{t-1} - \eta\beta\,\mathbf{v}_{t-1}$ rather than the current parameters:

$$\mathbf{v}_t = \beta\,\mathbf{v}_{t-1} + \nabla_{\boldsymbol{\theta}}\mathcal{L}\left(\boldsymbol{\theta}_{t-1} - \eta\beta\,\mathbf{v}_{t-1}\right)$$

If momentum is about to overshoot the basin, the lookahead gradient detects the opposing slope early and applies braking force. Scikit-learn’s MLPClassifier enables this by default (nesterovs_momentum=True).

AdaGrad and RMSProp: Per-Parameter Learning Rates

AdaGrad (Duchi, Hazan, Singer, 2011) accumulates squared gradients per parameter, dividing parameter updates by the square root of historical sum:

$$\begin{aligned} \mathbf{s}_t &= \mathbf{s}_{t-1} + \mathbf{g}_t^2 \\ \boldsymbol{\theta}_t &= \boldsymbol{\theta}_{t-1} - \frac{\eta}{\sqrt{\mathbf{s}_t} + \varepsilon}\,\mathbf{g}_t \end{aligned}$$

The scalar $\varepsilon \approx 10^{-8}$ prevents division by zero. Parameter $j$ receives effective learning rate $\eta / (\sqrt{s_j} + \varepsilon)$: historical steepness scales steps down, while rarely occurring sparse features receive larger steps. However, accumulator $\mathbf{s}_t$ strictly increases monotonically, eventually causing effective learning rates to decay toward zero and freezing training before convergence.

RMSProp (Hinton, 2012) resolves AdaGrad’s premature freezing by replacing cumulative sums with an exponential moving average, allowing older gradients to be forgotten:

$$\begin{aligned} \mathbf{s}_t &= \rho\,\mathbf{s}_{t-1} + (1 - \rho)\,\mathbf{g}_t^2 \\ \boldsymbol{\theta}_t &= \boldsymbol{\theta}_{t-1} - \frac{\eta}{\sqrt{\mathbf{s}_t} + \varepsilon}\,\mathbf{g}_t \end{aligned}$$

The decay rate $\rho$ is typically 0.9 or 0.99 (PyTorch default: 0.99). Here $\sqrt{\mathbf{s}_t}$ represents the root mean square (RMS) of recent gradients. If gradient magnitudes hover around $\lvert g\rvert$, then $\sqrt{s} \approx \lvert g\rvert$ and parameter step sizes evaluate to $\eta\,g / \lvert g\rvert$: roughly length $\eta$ regardless of raw gradient magnitude. At our toy ravine start, the gradient along $w_2$ is 19 times larger than $w_1$, yet RMSProp advances both dimensions by equal step lengths.

A key property: scaling the loss function by scalar $k$ multiplies gradients by $k$ and $\sqrt{\mathbf{s}_t}$ by $k$, leaving effective step sizes completely invariant. In standard GD, steps scale by $k$, breaking stability.

At step 1, because $\mathbf{s}_0 = \mathbf{0}$, $\mathbf{s}_1 = (1 - \rho)\,\mathbf{g}_1^2$, giving $\sqrt{\mathbf{s}_1} = \sqrt{1 - \rho}\,\lvert\mathbf{g}_1\rvert$. With $\rho = 0.99$, this evaluates to $0.1\,\lvert\mathbf{g}_1\rvert$, inflating the initial step length to 10 times $\eta$. Adam addresses this start-up artifact directly.

Adam: Combining Momentum with Adaptive Learning Rates

Adam (Adaptive Moment Estimation; Kingma and Ba, 2015) maintains two concurrent exponential moving averages: first moments (like Momentum) and uncentered second moments (like RMSProp):

$$\begin{aligned} \mathbf{v}_t &= \beta_1\,\mathbf{v}_{t-1} + (1 - \beta_1)\,\mathbf{g}_t \\ \mathbf{s}_t &= \beta_2\,\mathbf{s}_{t-1} + (1 - \beta_2)\,\mathbf{g}_t^2 \end{aligned}$$

Vector $\mathbf{v}_t$ estimates the expected gradient (first moment), while $\mathbf{s}_t$ estimates the expected squared gradient (second moment). Because Adam tracks two state vectors, it requires auxiliary memory twice the parameter count. Both moments are corrected for zero-initialization bias:

$$\begin{aligned} \hat{\mathbf{v}}_t &= \frac{\mathbf{v}_t}{1 - \beta_1^t}, \qquad \hat{\mathbf{s}}_t = \frac{\mathbf{s}_t}{1 - \beta_2^t} \\ \boldsymbol{\theta}_t &= \boldsymbol{\theta}_{t-1} - \eta\,\frac{\hat{\mathbf{v}}_t}{\sqrt{\hat{\mathbf{s}}_t} + \varepsilon} \end{aligned}$$

The exponent $\beta^t$ denotes the power of time step $t$. Default hyperparameters across the original paper, PyTorch, and scikit-learn are: $\eta = 0.001$, $\beta_1 = 0.9$, $\beta_2 = 0.999$, $\varepsilon = 10^{-8}$.

Why bias correction is essential. Because moving averages initialize at $\mathbf{0}$, early estimates are severely biased toward zero. At step 1 on our ravine with $\mathbf{g}_1 = [-8, 150]^\top$:

$$\begin{aligned} \mathbf{v}_1 &= 0.1\,\mathbf{g}_1 = [-0.8,\ 15]^\top \\ \mathbf{s}_1 &= 0.001\,\mathbf{g}_1^2 = [0.064,\ 22.5]^\top, \qquad \sqrt{\mathbf{s}_1} \approx [0.253,\ 4.743]^\top \end{aligned}$$

Uncorrected $\mathbf{v}_1$ is only 10% of true gradient, while $\sqrt{\mathbf{s}_1}$ is shrunk to $\sqrt{0.001} \approx 3.2\%$ of true magnitude. Without correction, their ratio evaluates to $[-3.16, 3.16]^\top$: the denominator is suppressed far more aggressively than the numerator, blowing up the first step to $0.1 / \sqrt{0.001} \approx 3.16\,\eta$. After bias correction, $\hat{\mathbf{v}}_1 = \mathbf{g}_1$ and $\hat{\mathbf{s}}_1 = \mathbf{g}_1^2$, restoring unit ratio $[-1, 1]^\top$ so each parameter steps exactly $\eta$.

Why dividing by $1 - \beta^t$ corrects bias. Assuming gradient remains constant $\mathbf{g}_t = \mathbf{g}$, unrolling the recursion:

$$\begin{aligned} \mathbf{v}_t &= (1 - \beta_1)\left(1 + \beta_1 + \dots + \beta_1^{t-1}\right)\mathbf{g} \\ &= (1 - \beta_1)\,\frac{1 - \beta_1^t}{1 - \beta_1}\,\mathbf{g} = \left(1 - \beta_1^t\right)\mathbf{g} \end{aligned}$$

The raw average is scaled by $(1 - \beta_1^t)$; dividing by this term recovers unbiased expectation $\mathbf{g}$. While $\beta_1 = 0.9$ dissipates its bias within tens of steps, $\beta_2 = 0.999$ recovers slowly: $1 - 0.999^{100} \approx 0.095$, meaning uncorrected $\mathbf{s}_{100}$ is only 10% of true scale after 100 steps. Without correction, early steps remain amplified for hundreds of iterations.

On our ravine, Adam ($\eta = 0.1$) converges in 193 steps—faster than GD (369), but behind Momentum (111) and RMSProp (103). Adam is not universally the fastest optimizer; its advantage lies in exceptional stability across diverse gradient scales.

Key Takeaway

Momentum accumulates gradients to accelerate along consistent directions. RMSProp scales steps inversely by recent gradient RMS, maintaining step lengths near $\eta$ across disparate scales. Adam combines both, scaling by $1 - \beta^t$ to counteract zero-initialization bias.

Learning Rate Schedules

Under mini-batch noise, maintaining constant $\eta$ causes parameters to bounce around the minimum. Decaying $\eta$ over time stabilizes convergence:

• Step decay: Multiplies $\eta$ by 0.1 at fixed epoch milestones.
• Cosine decay: Decays $\eta$ from $\eta_{\max}$ to $\eta_{\min}$ along a cosine half-cycle—gentle at the start and finish, steepest in the middle.
• Warmup: Linearly ramps $\eta$ from near zero to $\eta_{\max}$ during initial steps before decaying. This prevents unstable updates when parameter weights and optimizer moments are uncalibrated. Transformer and LLM training universally pair warmup with cosine decay.

Over total steps $T$, cosine decay follows:

$$\eta_t = \eta_{\min} + \frac{1}{2}\left(\eta_{\max} - \eta_{\min}\right)\left(1 + \cos\frac{\pi t}{T}\right)$$

At $t = 0$, $\eta_0 = \eta_{\max}$; at $t = T$, $\eta_T = \eta_{\min}$.

Practical Recommendations

Setting	Recommended Starting Point
Deep neural networks, new problems	Adam ($\eta = 10^{-3}$). Most forgiving baseline
Transformers and LLMs	AdamW with warmup and cosine decay
Computer vision CNNs (tuned)	SGD with Momentum 0.9 (or Nesterov) and step/cosine decay
Small linear models, tabular data	Standard solvers (Normal equation, L-BFGS in scikit-learn)

AdamW (Loshchilov and Hutter, 2019) decouples weight decay from gradient updates. In standard Adam with L2 penalty, regularizing gradients are divided by $\sqrt{\hat{\mathbf{s}}_t}$, inadvertently penalizing large-gradient weights less. AdamW subtracts decay term $\eta\lambda\boldsymbol{\theta}_{t-1}$ directly outside the adaptive fraction, ensuring uniform shrinkage across all parameters.

Limitations

• No single optimizer dominates all benchmarks. On simple quadratic surfaces, Momentum and RMSProp often outperform Adam. Wilson et al. (2017) observed that well-tuned SGD with momentum can yield better test generalization than adaptive methods.
• First-order limitations. All four methods remain strictly first-order gradient algorithms, mitigating but not eliminating condition number bottlenecks. Feature scaling remains essential.
• Non-convex landscapes. Complex networks feature non-convex saddle points and local minima where global convergence guarantees do not hold.

Visualization

The diagram below plots loss contours for $\mathcal{L} = \frac{1}{2}\left(w_1^2 + 100\,w_2^2\right)$ and the first 60 optimization steps starting from $(-8, 1.5)$:

The second code cell animates this comparison across all four optimizers alongside log-loss trajectories.

Code Examples

The three cells below run independently.

The first cell implements all four optimizers from scratch in NumPy, optimizing the toy ravine from $(-8, 1.5)$ until reaching a 0.01 error radius or 1,000 steps. A second run scales the loss by 10 (scale = 10) holding all learning rates constant.

Four optimizers on a narrow ravine

import numpy as np

c = np.array([1.0, 100.0])  # curvature along w1 and w2: a narrow valley


def run(name, eta, scale=1.0, beta1=0.9, beta2=0.999, rho=0.99, eps=1e-8):
    w = np.array([-8.0, 1.5])  # starting point
    v, s = np.zeros(2), np.zeros(2)  # first and second moment
    for t in range(1, 1001):
        g = scale * c * w  # gradient of scale * 0.5 * (w1^2 + 100 w2^2)
        if name == "GD":
            w = w - eta * g
        elif name == "Momentum":
            v = beta1 * v + g  # velocity: accumulates past gradients
            w = w - eta * v
        elif name == "RMSProp":
            s = rho * s + (1 - rho) * g**2  # mean of squares
            w = w - eta * g / (np.sqrt(s) + eps)
        else:  # Adam
            v = beta1 * v + (1 - beta1) * g
            s = beta2 * s + (1 - beta2) * g**2
            v_hat = v / (1 - beta1**t)  # bias correction
            s_hat = s / (1 - beta2**t)
            w = w - eta * v_hat / (np.sqrt(s_hat) + eps)
        dist = np.linalg.norm(w)  # distance to the minimum (0, 0)
        if dist < 0.01 or dist > 1e6:
            break
    status = "diverged" if dist > 1e6 else f"{t:4d} steps"
    return status, scale * 0.5 * np.sum(c * w**2)  # L when stopping


settings = [("GD", 0.018), ("Momentum", 0.018),
            ("RMSProp", 0.1), ("Adam", 0.1)]
for scale in [1.0, 10.0]:  # scale = 10: valley walls 10 times steeper
    print(f"scale = {scale:.0f}")
    for name, eta in settings:
        status, loss = run(name, eta, scale)
        print(f"  {name:8s} eta = {eta}: {status}, L = {loss:.1e}")

RunEdit code

The first run downloads Python, about 10 MB.

With scale = 1, iteration counts match theory: GD requires 369, Momentum 111, RMSProp 103, and Adam 193 steps. When scale = 10, curvature along $w_2$ inflates to 1000, reducing GD’s stability threshold to $2/1000 = 0.002$; consequently, $\eta = 0.018$ causes both GD and Momentum to diverge. In contrast, RMSProp and Adam converge in the exact same number of steps, demonstrating scale invariance.

The second cell plots 100-step trajectory paths across contour ellipses (left) alongside log-loss convergence curves (right).

Optimization paths on contour plots

import matplotlib.pyplot as plt
import numpy as np

c = np.array([1.0, 100.0])  # L = 0.5 * (w1^2 + 100 w2^2), as above


def path(name, eta, steps=100, beta1=0.9, beta2=0.999, rho=0.99):
    w, v, s = np.array([-8.0, 1.5]), np.zeros(2), np.zeros(2)
    points = [w]
    for t in range(1, steps + 1):
        g = c * w
        if name == "GD":
            step = g
        elif name == "Momentum":
            v = beta1 * v + g
            step = v
        elif name == "RMSProp":
            s = rho * s + (1 - rho) * g**2
            step = g / (np.sqrt(s) + 1e-8)
        else:  # Adam, with bias correction
            v = beta1 * v + (1 - beta1) * g
            s = beta2 * s + (1 - beta2) * g**2
            step = (v / (1 - beta1**t)) / (np.sqrt(s / (1 - beta2**t)) + 1e-8)
        w = w - eta * step
        points.append(w)
    return np.array(points)


fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4))
W1, W2 = np.meshgrid(np.linspace(-9, 3, 300), np.linspace(-2, 2, 300))
levels = np.geomspace(0.05, 150, 12)  # levels of L, evenly spaced in log
ax1.contour(W1, W2, 0.5 * (W1**2 + 100 * W2**2), levels, colors="0.8")
for name, eta in [("GD", 0.018), ("Momentum", 0.018),
                  ("RMSProp", 0.1), ("Adam", 0.1)]:
    p = path(name, eta)
    loss = 0.5 * (p[:, 0]**2 + 100 * p[:, 1]**2)  # L at each step
    ax1.plot(p[:, 0], p[:, 1], ".-", ms=3, lw=1, label=name)
    ax2.semilogy(loss, label=name)
    print(f"{name:8s}: L after 100 steps = {loss[-1]:.1e}")
ax1.plot(0, 0, "kx", ms=10)  # minimum
ax1.set(xlabel="w1", ylabel="w2", title="Paths, 100 steps")
ax2.set(xlabel="step", ylabel="loss", title="Loss (log scale)")
ax1.legend(loc="lower left")
plt.show()

RunEdit code

The first run downloads Python, about 10 MB.

After 100 iterations, RMSProp loss drops to $7.1 \times 10^{-5}$, Momentum reaches $1.1 \times 10^{-3}$, Adam sits at $0.45$, and GD at $0.85$. Momentum’s loss curve exhibits ripples from transversal bouncing. Adam smoothly hugs the $w_1$ axis without zigzagging. Note RMSProp’s large step-1 jump (length 1.0 rather than 0.1), confirming our theoretical analysis.

The third cell trains an MLPClassifier neural network (32 hidden neurons, ~2,400 parameters) on load_digits (1,797 handwritten $8 \times 8$ digits) across 20 epochs with batch size 64 and $\eta = 0.01$, comparing SGD, Momentum, and Adam.

SGD, Momentum, and Adam on load_digits

import warnings

import matplotlib.pyplot as plt
from sklearn.datasets import load_digits
from sklearn.exceptions import ConvergenceWarning
from sklearn.model_selection import train_test_split
from sklearn.neural_network import MLPClassifier

# 20 epochs is deliberately few, so silence the "not converged" warning
warnings.filterwarnings("ignore", category=ConvergenceWarning)

X, y = load_digits(return_X_y=True)  # 8x8 digit images, pixels 0..16
X_train, X_test, y_train, y_test = train_test_split(
    X / 16, y, test_size=0.25, random_state=0)  # scale pixels to [0, 1]

optimizers = {
    "SGD": dict(solver="sgd", momentum=0.0),
    "Momentum": dict(solver="sgd", momentum=0.9, nesterovs_momentum=False),
    "Adam": dict(solver="adam"),  # beta1 = 0.9, beta2 = 0.999 by default
}
for name, options in optimizers.items():
    # same network, same initialisation, same eta = 0.01, 64 images per batch
    model = MLPClassifier(hidden_layer_sizes=(32,), batch_size=64,
                          learning_rate_init=0.01, max_iter=20,
                          random_state=0, **options)
    model.fit(X_train, y_train)
    plt.plot(range(1, 21), model.loss_curve_, marker=".", label=name)
    loss, acc = model.loss_curve_[-1], model.score(X_test, y_test)
    print(f"{name:8s}: loss = {loss:.4f}, test acc = {acc:.3f}")

plt.xlabel("epoch")
plt.ylabel("training loss")
plt.yscale("log")
plt.legend()
plt.show()

RunEdit code

The first run downloads Python, about 10 MB.

After 20 epochs: SGD reaches loss 1.4786 with test accuracy 0.680; Momentum achieves loss 0.1815 and test accuracy 0.947; Adam drives loss down to 0.0321 with test accuracy 0.971.

Common Pitfalls

Switching Optimizers Without Re-tuning the Learning Rate

Each optimizer interprets $\eta$ differently. Under Momentum ($\beta = 0.9$), effective step length scales to $\approx 10\,\eta\,\lVert\mathbf{g}\rVert$, whereas Adam maintains steps near $\eta$ regardless of gradient scale. When switching optimizers, always reset $\eta$ to standard defaults: $10^{-3}$ for Adam, and $10^{-2}$ to $10^{-1}$ for SGD with momentum.

Omitting Bias Correction or Indexing Time from Step Zero

Omitting bias correction causes early updates to be 2 to 6.5 times larger than $\eta$ (under standard hyperparameters), frequently destabilizing training at the outset. Indexing time from $t = 0$ evaluates $1 - \beta^0 = 0$, causing division by zero. Always index steps starting at $t = 1$.

Assuming Adam Replaces Feature Standardization

While Adam adapts learning rates per coordinate axis, it cannot untangle diagonal correlations between non-orthogonal features. Standardizing inputs as shown in Gradient Descent remains computationally inexpensive and universally recommended.

Quick Check

Question 1 of 4 You use Momentum with $\beta = 0.9$. On a constant linear slope, the gradient remains constant at $\mathbf{g}$ over many steps. In steady state, how far does $\boldsymbol{\theta}$ advance at each step?

$\eta\mathbf{g}$, identical to standard GD$0.9\,\eta\mathbf{g}$$10\,\eta\mathbf{g}$$0.1\,\eta\mathbf{g}$

Check

Question 2 of 4 You run Adam ($\beta_1 = 0.9$, $\beta_2 = 0.999$, negligible $\varepsilon$, with bias correction enabled). At step 1, the gradient is $0.002$ along $\theta_1$ and $500$ along $\theta_2$. How far does each parameter step?

Both step exactly $\eta$, in the direction opposite their respective gradients$\theta_2$ steps roughly 250,000 times farther than $\theta_1$Both step $0.1\,\eta$, because $1 - \beta_1 = 0.1$Both step roughly $3.16\,\eta$

Check

Question 3 of 4 You scale the loss function by a factor of 100 (e.g., changing units of $y$) while keeping a previously well-tuned $\eta$ unchanged. What is most likely to happen to standard GD and Adam?

Both diverge because gradients are 100 times largerBoth behave identically because the minimum coordinates are unchangedAdam diverges, while GD remains unaffectedGD steps 100 times farther, easily exceeding the $2/L$ stability threshold and diverging; Adam remains virtually unchanged

Check

Question 4 of 4 You train a model extensively with AdaGrad. Loss remains high and gradients are non-zero, yet parameters have virtually ceased updating. What is the most plausible explanation?

The model reached a global minimumAccumulator $\mathbf{s}_t$ only adds positive squared terms and never decays, causing effective learning rate $\eta / \sqrt{\mathbf{s}_t}$ to asymptotically vanish toward 0Bias correction made step sizes too short$\varepsilon$ is excessively large compared to gradients

Check

Summary and Further Reading

Key Takeaways:

• Standard GD crawls when condition number $\kappa$ is large because steep directions constrain $\eta < 2/L$. On our $\kappa = 100$ ravine, GD requires 369 iterations.
• Momentum maintains velocity $\mathbf{v}_t = \beta\mathbf{v}_{t-1} + \mathbf{g}_t$, accelerating consistent directions by $1 / (1 - \beta)$ while damping transversal oscillations, at the cost of inertial overshoot.
• RMSProp scales updates by the square root of exponentially decayed squared gradients, keeping step lengths near $\eta$ and conferring loss-scale invariance. AdaGrad accumulates sums indefinitely, freezing prematurely.
• Adam tracks both first and second moments, applying $1 - \beta^t$ bias correction to prevent initial step blowup. Recommended default: Adam (or AdamW) with $\eta = 10^{-3}$ and a learning rate schedule.

Further Reading:

• Vũ Hữu Tiệp (2017). Machine Learning cơ bản, Bài 8: Gradient Descent (phần 2/2). Clear Vietnamese introduction to Momentum and Nesterov with animated illustrations.
• Zhang, A., Lipton, Z. C., Li, M., Smola, A. J. (2023). Dive into Deep Learning, Chapter 12: Optimization Algorithms. Sections 12.6 through 12.11 provide derivations and code using matching notation.
• Goh, G. (2017). Why Momentum Really Works. Distill. Visual, interactive breakdown of momentum dynamics on quadratic surfaces.
• Ruder, S. (2016). An overview of gradient descent optimization algorithms. arXiv:1609.04747. Comprehensive survey of modern optimization methods.
• Original Papers: Kingma & Ba (2015), Adam: A Method for Stochastic Optimization; Duchi et al. (2011), Adaptive Subgradient Methods (AdaGrad); Hinton et al. (2012), RMSProp; Loshchilov & Hutter (2019), Decoupled Weight Decay Regularization (AdamW).

I have finished this lesson