Optimizers in Deep Learning

The Story That Explains Optimizers

In machine learning, an optimizer is the algorithm that adjusts a model's parameters (weights and biases) to minimize the loss function — the measure of how wrong the model's predictions are. The process is called training, and every optimizer approaches it differently.

The Loss Landscape — A Visual Introduction

Before diving into individual optimizers, it helps to understand what they are navigating. The loss landscape is a high-dimensional surface where each point represents a configuration of model weights and its corresponding loss value.

SGD — Stochastic Gradient Descent

The Update Rule

SGD with Momentum — The Update Rule

Python Implementation

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

# Define a simple model
model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 10)
)

# SGD — plain vanilla
optimizer_sgd = optim.SGD(
    model.parameters(),
    lr=0.01           # learning rate
)

# SGD with momentum — almost always better
optimizer_sgd_m = optim.SGD(
    model.parameters(),
    lr=0.01,
    momentum=0.9,     # β — accumulates velocity
    weight_decay=1e-4 # L2 regularization
)

# SGD with Nesterov momentum (lookahead variant — often better)
optimizer_nag = optim.SGD(
    model.parameters(),
    lr=0.01,
    momentum=0.9,
    nesterov=True     # Nesterov Accelerated Gradient
)

# Training loop
criterion = nn.CrossEntropyLoss()

def train_step(X_batch, y_batch, optimizer):
    optimizer.zero_grad()           # clear old gradients
    output = model(X_batch)
    loss = criterion(output, y_batch)
    loss.backward()                  # compute gradients
    optimizer.step()                 # update weights
    return loss.item()

Adagrad — Adaptive Gradient Algorithm

The Update Rule

Python Implementation

import torch.optim as optim

# Adagrad — great for sparse problems (NLP, recommendation)
optimizer_adagrad = optim.Adagrad(
    model.parameters(),
    lr=0.01,             # global learning rate
    lr_decay=0,           # optional extra decay
    eps=1e-8,             # numerical stability
    weight_decay=0        # L2 regularization
)

# From scratch — educational implementation
import numpy as np

def adagrad_update(theta, grad, G, lr=0.01, eps=1e-8):
    """
    theta : current parameters (array)
    grad  : gradient of loss w.r.t theta
    G     : accumulated sum of squared gradients
    """
    G = G + grad**2                       # accumulate squared grads
    effective_lr = lr / (np.sqrt(G) + eps)  # per-param learning rate
    theta = theta - effective_lr * grad      # update parameters
    return theta, G

# Example with a sparse gradient scenario
theta = np.array([1.0, 1.0, 1.0])
G     = np.zeros_like(theta)           # starts at zero

# Feature 0 seen frequently, feature 2 rarely
grads = [
    np.array([0.5, 0.0, 0.0]),
    np.array([0.5, 0.0, 0.5]),
    np.array([0.5, 0.3, 0.0]),
    np.array([0.5, 0.0, 0.0]),
]

for i, grad in enumerate(grads):
    theta, G = adagrad_update(theta, grad, G)
    print(f"Step {i+1}: θ={theta.round(4)}, eff_lr={(0.01/(np.sqrt(G)+1e-8)).round(4)}")

RMSProp — Root Mean Square Propagation

The Update Rule

Python Implementation

import torch.optim as optim

# RMSProp — excellent for RNNs and RL
optimizer_rms = optim.RMSprop(
    model.parameters(),
    lr=0.001,            # learning rate
    alpha=0.99,           # β — decay for moving average (0.9 or 0.99)
    eps=1e-8,             # numerical stability
    weight_decay=0,       # L2 regularization
    momentum=0,           # optional momentum term
    centered=False         # if True: normalize by E[g] - E[g]² (centered variant)
)

# From scratch — numpy educational version
import numpy as np

def rmsprop_update(theta, grad, eg2, lr=0.001, beta=0.9, eps=1e-8):
    """
    theta : current parameters
    grad  : gradient of loss
    eg2   : running EWMA of squared gradients
    beta  : decay factor (0.9 = 10% new info per step)
    """
    eg2   = beta * eg2 + (1 - beta) * grad**2   # EWMA update
    theta = theta - lr / (np.sqrt(eg2) + eps) * grad  # adaptive update
    return theta, eg2

# Apply to a simple 1D loss: L(θ) = θ²  (minimum at θ=0)
theta = 5.0
eg2   = 0.0

for step in range(1, 21):
    grad       = 2 * theta                          # ∇(θ²) = 2θ
    theta, eg2 = rmsprop_update(theta, grad, eg2)
    if step % 5 == 0:
        print(f"Step {step:2d}: θ = {theta:.5f}, E[g²] = {eg2:.5f}")

Adam — Adaptive Moment Estimation

The Update Rule — All 4 Steps

Python Implementation

import torch.optim as optim

# Adam — the default choice for most deep learning
optimizer_adam = optim.Adam(
    model.parameters(),
    lr=3e-4,              # Andrej Karpathy's "learning rate 3e-4 is king"
    betas=(0.9, 0.999),   # β₁, β₂ — almost never tuned
    eps=1e-8,             # ε — numerical stability
    weight_decay=0,       # L2 regularization (see AdamW below)
    amsgrad=False          # AMSGrad variant (more stable, rarely needed)
)

# AdamW — Adam with CORRECT weight decay (strongly preferred)
optimizer_adamw = optim.AdamW(
    model.parameters(),
    lr=3e-4,
    betas=(0.9, 0.999),
    weight_decay=0.01      # decoupled weight decay — correct L2
)

# From scratch — full educational implementation
import numpy as np

class Adam:
    def __init__(self, params, lr=3e-4, betas=(0.9, 0.999), eps=1e-8):
        self.params = params
        self.lr     = lr
        self.b1, self.b2 = betas
        self.eps    = eps
        self.t      = 0
        self.m      = [np.zeros_like(p) for p in params]  # 1st moment
        self.v      = [np.zeros_like(p) for p in params]  # 2nd moment

    def step(self, grads):
        self.t += 1
        updated = []
        for i, (p, g) in enumerate(zip(self.params, grads)):
            # Step 1: first moment
            self.m[i] = self.b1 * self.m[i] + (1 - self.b1) * g
            # Step 2: second moment
            self.v[i] = self.b2 * self.v[i] + (1 - self.b2) * g**2
            # Step 3: bias correction
            m_hat = self.m[i] / (1 - self.b1**self.t)
            v_hat = self.v[i] / (1 - self.b2**self.t)
            # Step 4: update
            p_new = p - self.lr * m_hat / (np.sqrt(v_hat) + self.eps)
            updated.append(p_new)
        return updated

# Demo: minimise L(θ₁,θ₂) = θ₁² + 10·θ₂² (ill-conditioned)
params = [np.array(5.0), np.array(5.0)]
adam   = Adam(params, lr=0.1)

for step in range(1, 31):
    grads  = [2*params[0], 20*params[1]]    # ∇L
    params = adam.step(grads)
    if step % 10 == 0:
        loss = params[0]**2 + 10*params[1]**2
        print(f"Step {step:2d}: θ₁={params[0]:.4f}, θ₂={params[1]:.4f}, L={loss:.4f}")

Side-by-Side Comparison

Complete Python Implementation — All Four Optimizers

Full Training Loop Comparison on MNIST

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# ── Data ─────────────────────────────────────────────────────
transform = transforms.ToTensor()
train_data = datasets.MNIST('.', train=True,  download=True, transform=transform)
test_data  = datasets.MNIST('.', train=False, download=True, transform=transform)
train_loader = DataLoader(train_data, batch_size=128, shuffle=True)
test_loader  = DataLoader(test_data,  batch_size=256)

# ── Model ─────────────────────────────────────────────────────
def build_model():
    return nn.Sequential(
        nn.Flatten(),
        nn.Linear(784, 256), nn.ReLU(),
        nn.Linear(256, 128), nn.ReLU(),
        nn.Linear(128, 10)
    )

# ── Optimizers to compare ─────────────────────────────────────
configs = {
    'SGD':      lambda m: optim.SGD(m.parameters(), lr=0.01, momentum=0.9),
    'Adagrad':  lambda m: optim.Adagrad(m.parameters(), lr=0.01),
    'RMSProp':  lambda m: optim.RMSprop(m.parameters(), lr=0.001),
    'Adam':     lambda m: optim.Adam(m.parameters(),    lr=3e-4),
}

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
criterion = nn.CrossEntropyLoss()

def train_and_evaluate(name, optimizer_fn, epochs=5):
    model     = build_model().to(device)
    optimizer = optimizer_fn(model)
    history   = []

    for epoch in range(epochs):
        model.train()
        total_loss = 0

        for X, y in train_loader:
            X, y = X.to(device), y.to(device)
            optimizer.zero_grad()
            loss = criterion(model(X), y)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        # Evaluate
        model.eval()
        correct = 0
        with torch.no_grad():
            for X, y in test_loader:
                X, y = X.to(device), y.to(device)
                preds   = model(X).argmax(dim=1)
                correct += (preds == y).sum().item()

        acc = correct / len(test_data)
        history.append(acc)
        print(f"[{name}] Epoch {epoch+1}/{epochs} | Loss={total_loss/len(train_loader):.4f} | Test Acc={acc:.4f}")

    return history

# Run all four
results = {}
for name, opt_fn in configs.items():
    print(f"\n{'='*50}")
    results[name] = train_and_evaluate(name, opt_fn)

Learning Rate Schedules with Optimizers

# Best practice: combine Adam with a learning rate schedule
model     = build_model().to(device)
optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2)

# Option 1: Cosine Annealing — smoothly decays LR like a cosine wave
scheduler_cos = optim.lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=50, eta_min=1e-6
)

# Option 2: OneCycleLR — fast warmup then decay (Leslie Smith's policy)
scheduler_one = optim.lr_scheduler.OneCycleLR(
    optimizer,
    max_lr=1e-2,
    steps_per_epoch=len(train_loader),
    epochs=20
)

# Training loop with scheduler
for epoch in range(20):
    model.train()
    for X, y in train_loader:
        X, y = X.to(device), y.to(device)
        optimizer.zero_grad()
        loss = criterion(model(X), y)
        loss.backward()
        optimizer.step()
        scheduler_one.step()   # step INSIDE the batch loop for OneCycleLR

    current_lr = optimizer.param_groups[0]['lr']
    print(f"Epoch {epoch+1:2d} | LR = {current_lr:.6f}")

Hyperparameter Reference

When to Use Which Optimizer

Golden Rules