Chapter 8: Optimization

Bloom's Taxonomy Map for This Chapter

Learning Objectives

By the end of this chapter, you will be able to:

• Distinguish between Batch Gradient Descent, Stochastic Gradient Descent, and Mini-batch SGD — and quantify their trade-offs in speed, memory, and convergence stability
• Derive the exponentially weighted moving average formula and explain bias correction with numerical examples
• Explain how Momentum accelerates gradient descent using a rolling-ball analogy and implement it from scratch
• Describe how RMSprop adapts learning rates per-parameter to handle ill-conditioned loss surfaces
• State the complete Adam update equations with bias correction and justify the default hyperparameters (β₁=0.9, β₂=0.999, ε=10⁻⁸)
• Implement all five optimizers (SGD, Momentum, RMSprop, Adam, AdaGrad) as Python classes from scratch
• Compare optimizer loss curves on a common benchmark to evaluate convergence speed and stability
• Select appropriate learning rate schedules (step decay, exponential, cosine annealing, warm-up) for production training
• Design a training pipeline for distributed mini-batch SGD with learning rate warm-up, as used at Flipkart and similar Indian tech companies

Opening Hook — When 6 Hours Became 12 Minutes

Core Concepts — The Optimization Toolkit

Training a neural network means finding parameters W and b that minimize the loss function J(W, b). Gradient descent is the engine that drives this search, but how you compute and apply those gradients makes all the difference between a model that trains in minutes versus days. This section builds your toolkit from the ground up.

8.1 Gradient Descent Variants

All gradient descent algorithms share the same fundamental update rule:

The key difference between variants lies in how many samples you use to compute ∇J(θ) before each update.

Batch Size Trade-offs: A Complete View

8.2 Exponentially Weighted Averages (EWA)

Before we dive into Momentum and Adam, we need to understand the mathematical primitive they're built on: exponentially weighted moving averages. This is the single most important building block for all advanced optimizers.

Intuition: Smoothing Noisy Data

Imagine you're tracking daily temperatures in Delhi over a year. The raw data is noisy — 32°C one day, 28°C the next, 35°C the day after. To see the underlying trend, you compute a running average that gives more weight to recent values and exponentially less weight to older values.

What Does β Control?

The parameter β determines how many past values effectively contribute to the average. A rough approximation: V_t averages over approximately 1/(1−β) previous values.

Numerical Example

Let's trace through with β = 0.9 and daily temperatures θ = [35, 33, 36, 34, 38]:

# Tracing EWA step-by-step
V₀ = 0
V₁ = 0.9 × 0   + 0.1 × 35 = 3.5    # Way too low! (bias problem)
V₂ = 0.9 × 3.5 + 0.1 × 33 = 6.45
V₃ = 0.9 × 6.45+ 0.1 × 36 = 9.41
V₄ = 0.9 × 9.41+ 0.1 × 34 = 11.87
V₅ = 0.9 × 11.87+0.1 × 38 = 14.48
# After ~10 steps, V converges to the true range

The Bias Problem & Correction

Notice V₁ = 3.5 when the actual temperature is 35°C! Since V₀ = 0, the early estimates are biased toward zero. The fix is bias correction:

8.3 Gradient Descent with Momentum

The Rolling Ball Analogy

Imagine placing a ball at the top of a hilly landscape (the loss surface). Vanilla gradient descent is like the ball moving only according to the local slope — at every point it forgets its previous direction. Momentum is like giving the ball mass — it accumulates velocity in directions of consistent gradient, and resists changing direction for oscillatory gradients.

8.4 RMSprop — Root Mean Square Propagation

Momentum addresses the direction problem (dampening oscillations). But what about the magnitude problem? In many loss surfaces, some parameters have very large gradients while others have tiny ones. A single global learning rate is a poor fit for all of them.

8.5 Adam — Adaptive Moment Estimation

Adam combines the best of both worlds: Momentum's velocity (first moment of gradients) + RMSprop's adaptive scaling (second moment of gradients), with bias correction for both.

8.6 Learning Rate Schedules

A fixed learning rate is rarely optimal throughout training. Early on, you want large steps to explore quickly. Later, you want small steps to fine-tune near the minimum. Learning rate schedules systematically reduce α during training.

From-Scratch Code — Implementing Every Optimizer in Python

Let's implement all five optimizers as Python classes. Each takes parameters and gradients and applies the update rule. We'll then compare them on a common problem.

4a. Vanilla SGD

Python
import numpy as np

class SGD:
    """Vanilla Stochastic Gradient Descent."""
    def __init__(self, lr=0.01):
        self.lr = lr

    def update(self, params, grads):
        """
        params: dict of parameter arrays {'W1': ..., 'b1': ..., ...}
        grads:  dict of gradient arrays  {'dW1': ..., 'db1': ..., ...}
        """
        for key in params:
            params[key] -= self.lr * grads['d' + key]
        return params

4b. SGD with Momentum

Python
class MomentumSGD:
    """SGD with Momentum — rolling ball optimization."""
    def __init__(self, lr=0.01, beta=0.9):
        self.lr = lr
        self.beta = beta
        self.velocity = {}

    def update(self, params, grads):
        for key in params:
            if key not in self.velocity:
                self.velocity[key] = np.zeros_like(params[key])

            # Update velocity: V = β·V + (1-β)·dW
            self.velocity[key] = (self.beta * self.velocity[key]
                                  + (1 - self.beta) * grads['d' + key])

            # Update parameter: W = W - α·V
            params[key] -= self.lr * self.velocity[key]
        return params

4c. RMSprop

Python
class RMSprop:
    """RMSprop — adaptive per-parameter learning rates."""
    def __init__(self, lr=0.001, beta=0.999, epsilon=1e-8):
        self.lr = lr
        self.beta = beta
        self.epsilon = epsilon
        self.cache = {}  # Squared gradient accumulator

    def update(self, params, grads):
        for key in params:
            if key not in self.cache:
                self.cache[key] = np.zeros_like(params[key])

            grad = grads['d' + key]

            # Update cache: S = β·S + (1-β)·dW²
            self.cache[key] = (self.beta * self.cache[key]
                               + (1 - self.beta) * grad ** 2)

            # Update parameter: W = W - α·dW / (√S + ε)
            params[key] -= (self.lr * grad
                            / (np.sqrt(self.cache[key]) + self.epsilon))
        return params

4d. Adam — Full Implementation with Bias Correction

Python
class Adam:
    """
    Adam optimizer — Adaptive Moment Estimation.
    Combines Momentum (first moment) + RMSprop (second moment)
    with bias correction for both.
    """
    def __init__(self, lr=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8):
        self.lr = lr
        self.beta1 = beta1
        self.beta2 = beta2
        self.epsilon = epsilon
        self.m = {}  # First moment (mean of gradients)
        self.v = {}  # Second moment (mean of squared gradients)
        self.t = 0   # Time step counter

    def update(self, params, grads):
        self.t += 1

        for key in params:
            if key not in self.m:
                self.m[key] = np.zeros_like(params[key])
                self.v[key] = np.zeros_like(params[key])

            grad = grads['d' + key]

            # Step 1: Update first moment estimate (Momentum)
            self.m[key] = (self.beta1 * self.m[key]
                           + (1 - self.beta1) * grad)

            # Step 2: Update second moment estimate (RMSprop)
            self.v[key] = (self.beta2 * self.v[key]
                           + (1 - self.beta2) * grad ** 2)

            # Step 3: Bias correction
            m_hat = self.m[key] / (1 - self.beta1 ** self.t)
            v_hat = self.v[key] / (1 - self.beta2 ** self.t)

            # Step 4: Update parameters
            params[key] -= (self.lr * m_hat
                            / (np.sqrt(v_hat) + self.epsilon))
        return params

4e. AdaGrad (Bonus — Historical Importance)

Python
class AdaGrad:
    """
    AdaGrad — Adaptive Gradient Algorithm (Duchi et al., 2011).
    Accumulates ALL past squared gradients (no decay).
    Problem: learning rate monotonically decreases → may stop learning.
    """
    def __init__(self, lr=0.01, epsilon=1e-8):
        self.lr = lr
        self.epsilon = epsilon
        self.cache = {}

    def update(self, params, grads):
        for key in params:
            if key not in self.cache:
                self.cache[key] = np.zeros_like(params[key])

            grad = grads['d' + key]

            # Accumulate squared gradients (NO decay — key difference)
            self.cache[key] += grad ** 2

            # Update: divide by sqrt of total accumulated squared grads
            params[key] -= (self.lr * grad
                            / (np.sqrt(self.cache[key]) + self.epsilon))
        return params

4f. Comparing All Optimizers on the Same Problem

Let's train a simple 2-layer neural network on a synthetic classification task and overlay the loss curves.

Python
import numpy as np
import matplotlib.pyplot as plt

# ─── Synthetic Dataset ─────────────────────────────
np.random.seed(42)
N = 1000  # samples
D = 10    # features
X = np.random.randn(D, N)
W_true = np.random.randn(D, 1) * 0.5
y = (X.T @ W_true + np.random.randn(N, 1) * 0.1).T  # shape (1, N)

# ─── Simple Linear Model: y_hat = W.T @ x + b ─────
def init_params():
    return {
        'W': np.random.randn(D, 1) * 0.01,
        'b': np.zeros((1, 1))
    }

def compute_loss_and_grads(params, X, y):
    # Forward pass
    y_hat = params['W'].T @ X + params['b']  # (1, N)
    m = X.shape[1]
    loss = np.mean((y_hat - y) ** 2)

    # Backward pass
    diff = y_hat - y  # (1, N)
    grads = {
        'dW': (2 / m) * (X @ diff.T),  # (D, 1)
        'db': (2 / m) * np.sum(diff, axis=1, keepdims=True)
    }
    return loss, grads

# ─── Train with each optimizer ─────────────────────
optimizers = {
    'Vanilla SGD':   SGD(lr=0.01),
    'Momentum':      MomentumSGD(lr=0.01, beta=0.9),
    'RMSprop':       RMSprop(lr=0.001),
    'Adam':          Adam(lr=0.001),
    'AdaGrad':       AdaGrad(lr=0.1),
}

epochs = 200
results = {}

for name, optimizer in optimizers.items():
    params = init_params()
    losses = []
    for epoch in range(epochs):
        loss, grads = compute_loss_and_grads(params, X, y)
        losses.append(loss)
        params = optimizer.update(params, grads)
    results[name] = losses
    print(f"{name:15s} → Final loss: {losses[-1]:.6f}")

# ─── Plot comparison ───────────────────────────────
plt.figure(figsize=(10, 6))
colors = ['#ef4444', '#f59e0b', '#22c55e', '#7c3aed', '#64748b']
for (name, losses), color in zip(results.items(), colors):
    plt.plot(losses, label=name, linewidth=2, color=color)

plt.xlabel('Epoch', fontsize=12)
plt.ylabel('MSE Loss', fontsize=12)
plt.title('Optimizer Comparison: Loss Curves', fontsize=14, fontweight='bold')
plt.legend(fontsize=10)
plt.yscale('log')
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('optimizer_comparison.png', dpi=150)
plt.show()

4g. Learning Rate Schedules — Implementation

Python
import math

class StepDecaySchedule:
    """Reduce LR by factor every N epochs."""
    def __init__(self, initial_lr, drop_rate=0.1, step_size=30):
        self.initial_lr = initial_lr
        self.drop_rate = drop_rate
        self.step_size = step_size

    def get_lr(self, epoch):
        return self.initial_lr * (self.drop_rate ** (epoch // self.step_size))


class ExponentialDecaySchedule:
    """Smooth exponential decay."""
    def __init__(self, initial_lr, decay_rate=0.96):
        self.initial_lr = initial_lr
        self.decay_rate = decay_rate

    def get_lr(self, epoch):
        return self.initial_lr * (self.decay_rate ** epoch)


class CosineAnnealingSchedule:
    """Cosine annealing from lr_max to lr_min."""
    def __init__(self, lr_max, lr_min=1e-6, total_epochs=100):
        self.lr_max = lr_max
        self.lr_min = lr_min
        self.total_epochs = total_epochs

    def get_lr(self, epoch):
        return (self.lr_min + 0.5 * (self.lr_max - self.lr_min)
                * (1 + math.cos(math.pi * epoch / self.total_epochs)))


class WarmupCosineSchedule:
    """Linear warm-up followed by cosine annealing."""
    def __init__(self, lr_max, warmup_epochs=5,
                 total_epochs=100, lr_min=1e-6):
        self.lr_max = lr_max
        self.warmup_epochs = warmup_epochs
        self.total_epochs = total_epochs
        self.lr_min = lr_min

    def get_lr(self, epoch):
        if epoch < self.warmup_epochs:
            # Linear warm-up
            return self.lr_max * (epoch + 1) / self.warmup_epochs
        else:
            # Cosine annealing
            progress = (epoch - self.warmup_epochs) / (
                self.total_epochs - self.warmup_epochs)
            return (self.lr_min + 0.5 * (self.lr_max - self.lr_min)
                    * (1 + math.cos(math.pi * progress)))


# ─── Visualize all schedules ───────────────────────
epochs = 100
schedules = {
    'Step Decay (÷10 @ 30,60,90)': StepDecaySchedule(0.01),
    'Exponential (γ=0.96)':        ExponentialDecaySchedule(0.01),
    'Cosine Annealing':            CosineAnnealingSchedule(0.01),
    'Warmup + Cosine':             WarmupCosineSchedule(0.01, warmup_epochs=5),
}

plt.figure(figsize=(10, 5))
for name, schedule in schedules.items():
    lrs = [schedule.get_lr(e) for e in range(epochs)]
    plt.plot(lrs, label=name, linewidth=2)
plt.xlabel('Epoch'); plt.ylabel('Learning Rate')
plt.title('Learning Rate Schedules Comparison')
plt.legend(); plt.grid(True, alpha=0.3)
plt.tight_layout(); plt.show()

Industry Code — PyTorch Optimizers in Production

In real projects, you never implement optimizers from scratch. PyTorch provides battle-tested, GPU-accelerated implementations. Here's how to use them:

5a. Using PyTorch Optimizers

Python
import torch
import torch.nn as nn
import torch.optim as optim

# ─── Define a simple model ──────────────────────────
model = nn.Sequential(
    nn.Linear(10, 64),
    nn.ReLU(),
    nn.Linear(64, 32),
    nn.ReLU(),
    nn.Linear(32, 1)
)

# ─── Pick your optimizer ────────────────────────────
# Option 1: Vanilla SGD
optimizer = optim.SGD(model.parameters(), lr=0.01)

# Option 2: SGD with Momentum
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# Option 3: RMSprop
optimizer = optim.RMSprop(model.parameters(), lr=0.001, alpha=0.99)

# Option 4: Adam (most common default)
optimizer = optim.Adam(model.parameters(), lr=0.001,
                       betas=(0.9, 0.999), eps=1e-8)

# Option 5: AdamW (recommended for production)
optimizer = optim.AdamW(model.parameters(), lr=0.001,
                        betas=(0.9, 0.999), weight_decay=0.01)

# ─── Training loop ──────────────────────────────────
criterion = nn.MSELoss()

for epoch in range(100):
    for X_batch, y_batch in dataloader:
        # Forward pass
        y_pred = model(X_batch)
        loss = criterion(y_pred, y_batch)

        # Backward pass
        optimizer.zero_grad()   # CRITICAL: reset gradients
        loss.backward()         # Compute gradients
        optimizer.step()        # Apply optimizer update

    if (epoch + 1) % 10 == 0:
        print(f"Epoch {epoch+1:3d} | Loss: {loss.item():.6f}")

5b. PyTorch Learning Rate Schedulers

Python
from torch.optim.lr_scheduler import (
    StepLR, ExponentialLR, CosineAnnealingLR,
    CosineAnnealingWarmRestarts, OneCycleLR
)

optimizer = optim.AdamW(model.parameters(), lr=0.001)

# Step decay: reduce by 10× every 30 epochs
scheduler = StepLR(optimizer, step_size=30, gamma=0.1)

# Cosine annealing
scheduler = CosineAnnealingLR(optimizer, T_max=100, eta_min=1e-6)

# Cosine annealing with warm restarts
scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=2)

# One-cycle policy (super-convergence) — popular for fast training
scheduler = OneCycleLR(optimizer, max_lr=0.01,
                       total_steps=1000, pct_start=0.3)

# ─── Use in training loop ───────────────────────────
for epoch in range(100):
    for X_batch, y_batch in dataloader:
        optimizer.zero_grad()
        loss = criterion(model(X_batch), y_batch)
        loss.backward()
        optimizer.step()
    scheduler.step()  # Update LR after each epoch
    print(f"Epoch {epoch+1} | LR: {scheduler.get_last_lr()[0]:.6f}")

Visual Diagrams — Understanding Optimizer Behavior

6a. Contour Plot: Optimizer Paths on a 2D Loss Surface

Imagine the loss function as a bowl-shaped valley (contour lines show equal-loss regions). Each optimizer takes a different path to the minimum:

6b. The Optimizer Family Tree

6c. Bias Correction Visualization

Worked Example — Adam Step-by-Step on a 2D Problem

Let's manually trace 3 iterations of Adam on a simple problem to build deep intuition. We'll minimize f(w) = w² starting from w = 5.0.

Setup

Hyperparameters: α = 0.1, β₁ = 0.9, β₂ = 0.999, ε = 10⁻⁸
Initialize: m₀ = 0, v₀ = 0, t = 0

Iteration 1 (t = 1)

Iteration 2 (t = 2)

Iteration 3 (t = 3)

Case Study — Flipkart Search Ranking: Distributed Optimization at Scale

Common Mistakes & Misconceptions

Comprehensive Optimizer Comparison

Exercises

Section A — Multiple Choice Questions (10)

Section B — Short Answer Questions (5)

Section C — Long Answer Questions (3)

Section D — Programming Exercises (2)

Section E — Mini-Project

Chapter Summary

References & Further Reading

Core Papers

1. Kingma, D.P. & Ba, J. (2014). "Adam: A Method for Stochastic Optimization." ICLR 2015. arXiv:1412.6980. The original Adam paper — 180,000+ citations.
2. Loshchilov, I. & Hutter, F. (2019). "Decoupled Weight Decay Regularization." ICLR 2019. arXiv:1711.05101. The AdamW paper fixing Adam's weight decay bug.
3. Duchi, J., Hazan, E. & Singer, Y. (2011). "Adaptive Subgradient Methods for Online Learning and Stochastic Optimization." JMLR, 12, 2121–2159. The AdaGrad paper.
4. Polyak, B.T. (1964). "Some methods of speeding up the convergence of iterative methods." USSR Computational Mathematics and Mathematical Physics, 4(5), 1–17. The original momentum paper.
5. Hinton, G. (2012). "Lecture 6e: RMSprop." Coursera Neural Networks for Machine Learning. Unpublished — introduced in a lecture.

Important Related Work

6. Keskar, N.S. et al. (2017). "On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima." ICLR 2017. Why large batches hurt generalization.
7. Goyal, P. et al. (2017). "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour." arXiv:1706.02677. Linear LR scaling rule for distributed training.
8. Smith, L.N. & Topin, N. (2018). "Super-Convergence: Very Fast Training Using Large Learning Rates." arXiv:1708.07120. The 1-cycle policy.
9. Loshchilov, I. & Hutter, F. (2017). "SGDR: Stochastic Gradient Descent with Warm Restarts." ICLR 2017. Cosine annealing with restarts.
10. Wilson, A.C. et al. (2017). "The Marginal Value of Adaptive Gradient Methods in Machine Learning." NeurIPS 2017. SGD can outperform Adam on generalization.

Textbooks

11. Goodfellow, I., Bengio, Y. & Courville, A. (2016). Deep Learning. MIT Press. Chapter 8: Optimization for Training Deep Models.
12. Ruder, S. (2016). "An overview of gradient descent optimization algorithms." arXiv:1609.04747. Excellent survey comparing all optimizers.

Indian Industry References

13. Flipkart Tech Blog. "Scaling Deep Learning for Search Ranking." tech.flipkart.com
14. Infosys Research. "AI-Powered Fraud Detection for Indian Banking." Infosys Knowledge Institute, 2023.
15. TCS Research. "Multilingual NLP for Indian Languages: Training Strategies." ACL 2023 Workshop on Language Technology for Equality.