Stochastic Gradient Descent & Adaptive Optimizers

Stochastic Gradient Descent

In batch gradient descent, computing the full gradient over the entire dataset is prohibitively expensive when datasets contain millions of examples. SGD solves this by approximating the gradient using a small random subset of data, enabling scalable training of large models on massive datasets. The key insight is that a noisy but cheap gradient estimate is often better than an exact but expensive one.

SGD Algorithm

The full batch gradient requires computing ∇f(x) = (1/n) Σᵢ ∇fᵢ(x) over all n samples. SGD replaces this with ∇f_{i_k}(x) for a single random index i_k, reducing per-iteration cost from O(n) to O(1).

Mini-Batch SGD

In practice, purely stochastic (batch size = 1) gradients are too noisy. Mini-batch SGD strikes a balance by sampling a small batch B of size b from the training set:

Common batch sizes in practice: 32, 64, 128, 256. Larger batches provide more accurate gradients but require more memory and may generalize worse.

Why SGD Works

Proof sketch: Since each data point i is sampled uniformly at random:

E[g_k | x_k] = E[(1/b) Σᵢ∈B ∇fᵢ(x_k)] = (1/n) Σᵢ₌₁ⁿ ∇fᵢ(x_k) = ∇f(x_k)

This unbiasedness guarantee ensures that, given a constant learning rate α and Lipschitz continuous gradients, SGD converges to a neighborhood of the optimum with radius proportional to α. Decreasing the learning rate over time shrinks this neighborhood, enabling convergence to the exact solution.

Variance of SGD Gradients

The variance σ² directly affects SGD convergence speed. High variance means noisier updates, requiring smaller learning rates and more iterations. This motivates variance reduction techniques.

The first term O(1/K) decreases with more iterations. The second term ασ²/2 is the "noise floor" — the irreducible error due to gradient variance. Reducing α shrinks the noise floor but slows convergence, creating a fundamental trade-off.

Variance Reduction

Several techniques reduce SGD variance without requiring full-batch gradients:

SVRG (Stochastic Variance Reduced Gradient): Periodically computes a full gradient to anchor the stochastic estimate, reducing variance to zero near the solution.

SAGA: Maintains a table of individual gradients, providing unbiased variance reduction with linear convergence for strongly convex functions.

Importance Sampling: Sampling data points with probability proportional to their gradient magnitude reduces variance compared to uniform sampling.

Momentum

SGD with momentum accelerates convergence by accumulating velocity in directions of consistent gradient.

Why momentum helps:

• Accumulates velocity in consistent gradient directions, accelerating convergence along ravines
• Dampens oscillations in directions with high curvature
• Effective learning rate is amplified by factor 1/(1-β) ≈ 10 for β=0.9

AdaGrad

Key properties:

• Parameters with large accumulated gradients get smaller effective learning rates
• Parameters with sparse gradients get larger effective learning rates
• Well-suited for sparse data (NLP, recommender systems)

Critical limitation: The accumulator G_k only grows, causing the learning rate to decay aggressively and eventually approach zero. This can halt training prematurely.

class AdaGrad:
    def __init__(self, lr=0.01, eps=1e-8):
        self.lr = lr
        self.eps = eps
        self.G = None

    def step(self, x, grad):
        if self.G is None:
            self.G = np.zeros_like(x)
        self.G += grad ** 2
        return x - self.lr * grad / (np.sqrt(self.G) + self.eps)

RMSProp

RMSProp vs AdaGrad: RMSProp replaces the growing accumulator with an exponential moving average, preventing the learning rate from decaying to zero. The effective window is approximately 1/(1-β) ≈ 10 recent gradients.

class RMSProp:
    def __init__(self, lr=0.001, beta=0.9, eps=1e-8):
        self.lr = lr
        self.beta = beta
        self.eps = eps
        self.v = None

    def step(self, x, grad):
        if self.v is None:
            self.v = np.zeros_like(x)
        self.v = self.beta * self.v + (1 - self.beta) * grad ** 2
        return x - self.lr * grad / (np.sqrt(self.v) + self.eps)

Adam Optimizer

Adam Full Algorithm

Algorithm: Adam Optimizer
---------------------------------------------
Require: Learning rate α (default: 0.001)
Require: Decay rates β₁ = 0.9, β₂ = 0.999
Require: Numerical stability ε = 1e-8
Require: Initial parameters x₀

1:  Initialize m₀ <- 0, v₀ <- 0, t <- 0
2:  repeat
3:      t <- t + 1
4:      Sample mini-batch B_k
5:      g_k <- (1/|B_k|) Σᵢ∈B_k ∇fᵢ(x_{t-1})
6:      m_t <- β₁ · m_{t-1} + (1 - β₁) · g_t       [Update biased first moment]
7:      v_t <- β₂ · v_{t-1} + (1 - β₂) · g_t²       [Update biased second moment]
8:      m̂_t <- m_t / (1 - β₁ᵗ)                       [Bias correction]
9:      v̂_t <- v_t / (1 - β₂ᵗ)                       [Bias correction]
10:     x_t <- x_{t-1} - α · m̂_t / (√v̂_t + ε)       [Parameter update]
11: until convergence
---------------------------------------------

class Adam:
    def __init__(self, lr=0.001, beta1=0.9, beta2=0.999, eps=1e-8):
        self.lr = lr
        self.beta1 = beta1
        self.beta2 = beta2
        self.eps = eps
        self.m = None
        self.v = None
        self.t = 0

    def step(self, x, grad):
        self.t += 1
        if self.m is None:
            self.m = np.zeros_like(x)
            self.v = np.zeros_like(x)

        self.m = self.beta1 * self.m + (1 - self.beta1) * grad
        self.v = self.beta2 * self.v + (1 - self.beta2) * grad ** 2

        m_hat = self.m / (1 - self.beta1 ** self.t)
        v_hat = self.v / (1 - self.beta2 ** self.t)

        return x - self.lr * m_hat / (np.sqrt(v_hat) + self.eps)

AdamW: Decoupled Weight Decay

Learning Rate Warmup

Why warmup is needed:

• At initialization, Adam's second moment estimates v̂ are noisy and unreliable
• Large gradients in early training can cause divergent updates
• Warmup allows moment estimates to stabilize before using full learning rates
• Especially important for transformers and large batch training

Common warmup schedules:

def warmup_lr(step, warmup_steps, target_lr):
    if step < warmup_steps:
        return target_lr * step / warmup_steps
    return target_lr

Python Implementation

import numpy as np

class SGD:
    def __init__(self, lr=0.01, momentum=0.0, weight_decay=0.0):
        self.lr = lr
        self.momentum = momentum
        self.weight_decay = weight_decay
        self.v = None

    def step(self, x, grad):
        if self.v is None:
            self.v = np.zeros_like(x)

        if self.weight_decay > 0:
            grad = grad + self.weight_decay * x

        self.v = self.momentum * self.v + grad
        return x - self.lr * self.v


class AdaGrad:
    def __init__(self, lr=0.01, eps=1e-8):
        self.lr = lr
        self.eps = eps
        self.G = None

    def step(self, x, grad):
        if self.G is None:
            self.G = np.zeros_like(x)
        self.G += grad ** 2
        return x - self.lr * grad / (np.sqrt(self.G) + self.eps)


class RMSProp:
    def __init__(self, lr=0.001, beta=0.9, eps=1e-8):
        self.lr = lr
        self.beta = beta
        self.eps = eps
        self.v = None

    def step(self, x, grad):
        if self.v is None:
            self.v = np.zeros_like(x)
        self.v = self.beta * self.v + (1 - self.beta) * grad ** 2
        return x - self.lr * grad / (np.sqrt(self.v) + self.eps)


class Adam:
    def __init__(self, lr=0.001, beta1=0.9, beta2=0.999, eps=1e-8):
        self.lr = lr
        self.beta1 = beta1
        self.beta2 = beta2
        self.eps = eps
        self.m = None
        self.v = None
        self.t = 0

    def step(self, x, grad):
        self.t += 1
        if self.m is None:
            self.m = np.zeros_like(x)
            self.v = np.zeros_like(x)

        self.m = self.beta1 * self.m + (1 - self.beta1) * grad
        self.v = self.beta2 * self.v + (1 - self.beta2) * grad ** 2

        m_hat = self.m / (1 - self.beta1 ** self.t)
        v_hat = self.v / (1 - self.beta2 ** self.t)

        return x - self.lr * m_hat / (np.sqrt(v_hat) + self.eps)


class AdamW:
    def __init__(self, lr=0.001, beta1=0.9, beta2=0.999, eps=1e-8, weight_decay=0.01):
        self.lr = lr
        self.beta1 = beta1
        self.beta2 = beta2
        self.eps = eps
        self.weight_decay = weight_decay
        self.m = None
        self.v = None
        self.t = 0

    def step(self, x, grad):
        self.t += 1
        if self.m is None:
            self.m = np.zeros_like(x)
            self.v = np.zeros_like(x)

        self.m = self.beta1 * self.m + (1 - self.beta1) * grad
        self.v = self.beta2 * self.v + (1 - self.beta2) * grad ** 2

        m_hat = self.m / (1 - self.beta1 ** self.t)
        v_hat = self.v / (1 - self.beta2 ** self.t)

        return x - self.lr * (m_hat / (np.sqrt(v_hat) + self.eps) + self.weight_decay * x)


# --- Demo: Compare optimizers on Rosenbrock function ---
def rosenbrock(x):
    return (1 - x[0])**2 + 100 * (x[1] - x[0]**2)**2

def rosenbrock_grad(x):
    dx0 = -2 * (1 - x[0]) + 200 * (x[1] - x[0]**2) * (-2 * x[0])
    dx1 = 200 * (x[1] - x[0]**2)
    return np.array([dx0, dx1])

optimizers = {
    "SGD": SGD(lr=0.001, momentum=0.9),
    "AdaGrad": AdaGrad(lr=0.05),
    "RMSProp": RMSProp(lr=0.001),
    "Adam": Adam(lr=0.005),
    "AdamW": AdamW(lr=0.005, weight_decay=0.01),
}

for name, opt in optimizers.items():
    x = np.array([-1.0, 1.0])
    for i in range(5000):
        grad = rosenbrock_grad(x)
        x = opt.step(x, grad)
    print(f"{name:10s} | Final loss: {rosenbrock(x):.6f} | Position: ({x[0]:.4f}, {x[1]:.4f})")

Applications in AI/ML

Deep Learning Training

SGD and its variants are used in virtually all neural network training:

Learning Rate Schedules in Practice

Cosine Annealing:

One Cycle Policy:

• Linearly warmup from low LR to peak LR for first ~30% of training
• Cosine decay from peak to very low LR for remaining ~70%
• Popularized by Leslie Smith; used in fast.ai course

Common Mistakes

Interview Questions

Q1: Why is SGD preferred over batch gradient descent for deep learning? A: Batch GD requires computing the gradient over the entire dataset per update, which is infeasible for millions of samples. SGD uses mini-batches, enabling faster iteration, lower memory usage, and the noise helps escape saddle points. The noise also acts as implicit regularization, often improving generalization.

Q2: What is the role of the learning rate in SGD, and how do you choose it? A: The learning rate controls step size. Too large causes divergence; too small causes slow convergence. In practice, use a learning rate finder (sweep log-scale LR and plot loss), start with Adam lr=1e-3 or SGD lr=0.1, and apply cosine annealing or step decay schedules.

Q3: Explain bias correction in Adam. Why is it necessary? A: At initialization m₀=v₀=0, so early EMA estimates are biased toward zero. Dividing by (1-βᵗ) corrects this. For β₁=0.9, the correction factor at t=1 is 1/0.1=10, which is critical. After ~20 steps, the correction becomes negligible (<5%).

Q4: Why does AdamW outperform Adam with L2 regularization? A: In Adam, L2 regularization adds λx to the gradient, which is then scaled by 1/√v̂. This means the effective weight decay varies per-parameter inversely to gradient magnitude. AdamW decouples weight decay, applying λx directly without adaptive scaling, providing uniform regularization.

Q5: When would you use SGD+momentum over Adam? A: SGD+momentum often achieves better final generalization on vision tasks (ImageNet, etc.) despite slower convergence. Adam converges faster but may converge to sharper minima. Common pattern: train with Adam for speed, then switch to SGD for fine-tuning.

Q6: What happens if you use β₂=0.999 with a very large batch size? A: The second moment estimate v̂ becomes a very long-running average, potentially stale for non-stationary objectives. Consider reducing β₂ to 0.99 or 0.95 for GANs or fast-changing loss landscapes.

Q7: Explain the relationship between batch size and learning rate. A: Linear scaling rule: when batch size is multiplied by k, multiply LR by k. This maintains the same signal-to-noise ratio in gradient updates. Must be combined with warmup to avoid early instability. (Goyal et al., 2017)

Practice Problems

Quick Reference

Default hyperparameters:

Key formulas:

• SGD: x <- x - α·g
• Momentum: v <- βv + g; x <- x - αv
• AdaGrad: G <- G + g²; x <- x - αg/√(G+ε)
• RMSProp: v <- 0.9v + 0.1g²; x <- x - αg/√(v+ε)
• Adam: m <- 0.9m + 0.1g; v <- 0.999v + 0.001g²; x <- x - α·m̂/(√v̂+ε)

Cross-References

• Gradient Descent (062): The deterministic foundation; batch gradient descent converges exactly but is slow.
• Newton's Method (064): Second-order methods that use curvature information for faster convergence.
• Convex Optimization (061): Theoretical convergence guarantees for SGD in convex settings.
• Constrained Optimization (065): Extends SGD to constrained problems via projected variants.
• Hyperparameter Optimization (069): Finding optimal learning rates and batch sizes systematically.
• Calculus: Partial Derivatives (027): Foundation for computing gradients in multi-dimensional optimization.
• Calculus: Optimization (030): Classical unconstrained optimization theory and optimality conditions.
• Linear Algebra: Norms (015): Measuring gradient magnitude and parameter distance in optimization.
• Probability: Expectation (039): Expected value analysis of SGD convergence.
• Probability: CLT (042): Central Limit Theorem explains why mini-batch gradients concentrate around the true gradient.