Deep Learning

Training Deep Networks Effectively

Optimizers: SGD, Momentum, RMSProp, Adam

Module 2, Chapter 5 implemented plain mini-batch gradient descent: W -= learning_rate * gradient. In practice, this converges slowly and struggles with two comm

JrCodex·6 min read

Jr Codex Deep Learning Notes

Level: Intermediate Prerequisites: Chapter 1: Weight Initialization Time to complete: ~25 minutes


Table of Contents

  1. The Limits of Plain Gradient Descent
  2. Momentum
  3. RMSProp
  4. Adam — Combining Both Ideas
  5. Comparing Optimizers Visually
  6. Optimizers in PyTorch
  7. Which Optimizer Should You Use?
  8. Summary & Next Steps

1. The Limits of Plain Gradient Descent

Module 2, Chapter 5 implemented plain mini-batch gradient descent: W -= learning_rate * gradient. In practice, this converges slowly and struggles with two common landscape shapes:

Two Problems Plain Gradient Descent Struggles With
─────────────────────────────────────────
  Ravines:      the loss surface is much steeper in one
                  direction than another — plain gradient
                  descent OSCILLATES across the steep direction
                  while making slow progress along the shallow one
  Saddle points:    flat regions where the gradient is nearly
                       zero in some directions but the point
                       isn't actually a minimum — plain gradient
                       descent can get stuck here for a long time
─────────────────────────────────────────

2. Momentum

Momentum addresses oscillation in ravines by accumulating a "velocity" — a running average of past gradients — so consistent gradient directions build up speed, while oscillating directions cancel out.

import numpy as np
 
def gradient_descent_with_momentum(gradients_fn, W_init, learning_rate=0.01, momentum=0.9, steps=100):
    W = W_init.copy()
    velocity = np.zeros_like(W)
 
    for step in range(steps):
        gradient = gradients_fn(W)
        velocity = momentum * velocity - learning_rate * gradient      # accumulate velocity
        W += velocity                                                     # move using velocity, not just the raw gradient
 
    return W
The Physical Analogy
─────────────────────────────────────────
  Plain gradient descent:    a ball that stops and re-aims
                                at every single step, based
                                purely on the CURRENT slope
  Momentum:                      a ball rolling downhill,
                                    building up speed in a
                                    CONSISTENT direction, and
                                    resisting sudden direction
                                    changes from noisy individual
                                    gradients
─────────────────────────────────────────

The momentum term (typically 0.9) controls how much of the previous velocity carries over — higher values produce more "inertia."


3. RMSProp

RMSProp takes a different approach: it adapts the learning rate per parameter, dividing by a running average of that parameter's recent gradient magnitudes — parameters with consistently large gradients get a smaller effective learning rate, and vice versa.

import numpy as np
 
def rmsprop(gradients_fn, W_init, learning_rate=0.01, decay_rate=0.9, epsilon=1e-8, steps=100):
    W = W_init.copy()
    squared_grad_avg = np.zeros_like(W)
 
    for step in range(steps):
        gradient = gradients_fn(W)
        squared_grad_avg = decay_rate * squared_grad_avg + (1 - decay_rate) * gradient**2
        W -= learning_rate * gradient / (np.sqrt(squared_grad_avg) + epsilon)      # PER-PARAMETER adaptive rate
 
    return W
Why Per-Parameter Adaptation Helps
─────────────────────────────────────────
  In a real network, DIFFERENT weights can have VERY
  different gradient scales — some parameters need large
  updates, others need tiny, careful updates. A SINGLE global
  learning rate (Module 2, Ch.5) is a compromise across ALL
  of them; RMSProp lets each parameter effectively have its
  OWN learning rate, based on its own gradient history.
─────────────────────────────────────────

4. Adam — Combining Both Ideas

Adam (Adaptive Moment Estimation) combines Momentum's velocity tracking with RMSProp's per-parameter adaptive rates — it is, by a wide margin, the most commonly used optimizer in modern deep learning, and the default choice for virtually every architecture in this curriculum from Module 4 onward.

import numpy as np
 
def adam(gradients_fn, W_init, learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8, steps=100):
    W = W_init.copy()
    m = np.zeros_like(W)      # momentum term (first moment)
    v = np.zeros_like(W)      # RMSProp term (second moment)
 
    for t in range(1, steps + 1):
        gradient = gradients_fn(W)
 
        m = beta1 * m + (1 - beta1) * gradient           # Momentum's velocity
        v = beta2 * v + (1 - beta2) * gradient**2            # RMSProp's squared gradient average
 
        # Bias correction — compensates for m, v starting at zero
        m_corrected = m / (1 - beta1**t)
        v_corrected = v / (1 - beta2**t)
 
        W -= learning_rate * m_corrected / (np.sqrt(v_corrected) + epsilon)
 
    return W
Why Adam Became the Default
─────────────────────────────────────────
  1. Combines momentum's smoothing with RMSProp's per-
       parameter adaptation — gets the benefits of BOTH
  2. Works well with LITTLE hyperparameter tuning — the
       default values (lr=0.001, beta1=0.9, beta2=0.999) work
       reasonably well across a WIDE range of problems
  3. The BIAS CORRECTION step specifically fixes a slow-start
       problem that would otherwise occur early in training,
       when m and v are still close to their zero
       initialization
─────────────────────────────────────────

5. Comparing Optimizers Visually

Convergence Behavior on a Typical Loss Surface (Conceptual)
─────────────────────────────────────────
  Plain SGD:       ╲╱╲╱╲╱╲___________
                     (oscillates, slow along shallow direction)

  SGD + Momentum:  ╲___________
                     (smoother, builds speed, less oscillation)

  RMSProp:         ╲__________
                     (adapts per-parameter, avoids ravines
                       differently than momentum)

  Adam:            ╲_______
                     (typically fastest, most reliable
                       convergence across problem types)
─────────────────────────────────────────

6. Optimizers in PyTorch

import torch
import torch.nn as nn
 
model = nn.Sequential(nn.Linear(10, 20), nn.ReLU(), nn.Linear(20, 1))
 
# Plain SGD (Module 2, Ch.5) — rarely used alone in practice anymore
optimizer_sgd = torch.optim.SGD(model.parameters(), lr=0.01)
 
# SGD with momentum (Section 2)
optimizer_momentum = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
 
# RMSProp (Section 3)
optimizer_rmsprop = torch.optim.RMSprop(model.parameters(), lr=0.001)
 
# Adam (Section 4) — the default choice for most of this curriculum
optimizer_adam = torch.optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.999))
 
# AdamW — Adam with IMPROVED weight decay (regularization, Chapter 4) — the
# modern default for training Transformers (Module 6) specifically
optimizer_adamw = torch.optim.AdamW(model.parameters(), lr=0.001, weight_decay=0.01)

7. Which Optimizer Should You Use?

Practical Guidance
─────────────────────────────────────────
  Default choice:         Adam (or AdamW) — works well across
                             the vast majority of architectures
                             and problems in this curriculum
  Computer vision                SGD + Momentum, with careful
  (sometimes):                      learning rate scheduling
                                      (Chapter 5), can slightly
                                      OUTPERFORM Adam for some
                                      CNN architectures (Module 4)
                                      — a known research finding,
                                      but Adam remains a strong
                                      and simpler default
  Transformers                       AdamW is the near-universal
  (Module 6):                           standard — used to train
                                          virtually every modern
                                          large language model
─────────────────────────────────────────

8. Summary & Next Steps

Key Takeaways

  • Plain gradient descent struggles with ravines (oscillation) and saddle points (stalling); Momentum and RMSProp each address one of these limitations.
  • Momentum accumulates a velocity from past gradients, smoothing oscillation and building speed in consistent directions.
  • RMSProp adapts the effective learning rate per parameter, based on each parameter's recent gradient magnitude history.
  • Adam combines both ideas and is the default optimizer for nearly every architecture in this curriculum; AdamW (Adam with improved weight decay) is the standard for training Transformers specifically.

Concept Check

  1. What specific loss-surface problem does Momentum address, and how?
  2. Why does RMSProp use a per-parameter learning rate instead of one global rate?
  3. What are the two ideas Adam combines, and why does it include a "bias correction" step?

Next Chapter

Chapter 3: Batch Normalization & Layer Normalization


Jr Codex — 1-on-1 Personalized Coaching | Back to Module Index | Back to DL Index