Skip to content

Optimization for ML

Training a model is optimization: minimize a loss over millions of parameters. This sheet covers gradient descent, modern optimizers, learning rates, regularization-as-constraint, and search.

Table of Contents

Installation / Setup

# General-purpose optimization (small problems):
pip install scipy            # scipy.optimize
# Deep learning optimizers:
pip install torch            # torch.optim

Core Syntax

The optimization problem

$$\theta^* = \arg\min_\theta \mathcal{L}(\theta) = \arg\min_\theta \frac{1}{n}\sum_{i=1}^{n} \ell(f(x_i; \theta), y_i)$$

Term Meaning
$\theta$ Parameters (weights, biases)
$\mathcal{L}$ Loss/objective (MSE, cross-entropy, ...)
$\ell$ Per-sample loss
$n$ Number of training examples

Gradient descent

$$\theta \leftarrow \theta - \eta \nabla_\theta \mathcal{L}(\theta)$$

  • $\eta$ (learning rate): step size. Direction comes from the gradient (see calculus.md).
  • Too large: overshoot/diverge. Too small: crawling. Roughly right: steady decrease.
w -= lr * grad              # the whole algorithm in one line

Batch variants

Variant Gradient estimated from Notes
Batch GD whole dataset exact but expensive
Stochastic GD 1 random sample noisy but cheap; can escape poor regions
Mini-batch SGD small batch (e.g. 32–256) the default in deep learning

SGD's noise often helps generalization; smaller batches add regularization-like effects.

Optimizer zoo

Optimizer Update idea When to use
SGD plain step baseline, CV-style models
SGD + momentum keep a velocity $v \leftarrow \mu v - \eta g$; $\theta \leftarrow \theta + v$ smoother, faster convergence
NAG look-ahead momentum slightly better than momentum
AdaGrad per-param scaling by accumulated squares sparse features
RMSProp moving-average scaling non-stationary, RNNs
Adam momentum + RMSProp + bias correction default for most DL
AdamW Adam with decoupled weight decay transformers; usually better than Adam

Adam defaults that usually work: $\eta = 10^{-3}$, $\beta_1 = 0.9$, $\beta_2 = 0.999$, $\epsilon = 10^{-8}$.

import torch.optim as optim
opt = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
# loop: opt.zero_grad(); loss.backward(); opt.step()

Learning-rate schedules

Schedule Behavior Use
Constant fixed $\eta$ simple baselines
Step decay halve every $k$ epochs classic CNNs
Cosine smooth decay to ~0 modern SOTA, transformers
Warmup small $\eta$ → peak → schedule stable big-batch / transformer training
OneCycle up then down fast convergence on vision

Rule of thumb: if loss explodes, lower $\eta$; if it plateaus early, schedule decay.

Convex vs non-convex

  • Convex (MSE/logistic regression, linear SVMs): every local minimum is global → fast, reliable.
  • Non-convex (neural networks): many minima and saddle points → initialization, schedules, and optimizers matter.
  • Check calculus.md for the Hessian-based convexity test.

Regularization as constrained optimization

Unconstrained form: $\min_\theta \mathcal{L}(\theta) + \lambda |\theta|_p^p$ is the Lagrangian of a constraint $|\theta|_p^p \le c$.

Penalty Effect ML equivalent
L2, $\lambda|\theta|_2^2$ shrinks weights smoothly ridge, weight decay
L1, $\lambda|\theta|_1$ drives weights to exactly 0 lasso, sparsity
  • Larger $\lambda$ = stronger constraint = more bias, less variance (see statistics.md).
  • Weight decay in AdamW ≈ L2 with better behavior for adaptive methods.

Constrained optimization (general)

Lagrange multipliers: minimize $f(x)$ subject to $g(x) = 0$ → solve $\nabla f = \lambda \nabla g$. KKT conditions generalize to inequalities — the theory behind SVMs (maximize margin subject to constraints) and dual problems.

Hyperparameter optimization

Method Idea Cost
Grid search all combos expensive
Random search sample ranges often better than grid
Bayesian (Optuna) learn a surrogate, suggest promising points efficient 
import optuna
def objective(trial):
    lr = trial.suggest_float("lr", 1e-4, 1e-1, log=True)
    ...
    return val_loss
study = optuna.create_study(direction="minimize")
study.optimize(objective, n_trials=50)

Common Operations

import numpy as np
from scipy.optimize import minimize

# Small, smooth problem -> L-BFGS (quasi-Newton) often best
res = minimize(lambda w: ((X @ w - y) ** 2).mean(), x0=np.zeros(3),
               method="L-BFGS-B")
res.x, res.success, res.message

# Custom numpy mini-batch GD
rng = np.random.default_rng(0)
w = np.zeros(X.shape[1])
for epoch in range(200):
    idx = rng.choice(len(X), size=64, replace=False)
    g = 2 * X[idx].T @ (X[idx] @ w - y[idx]) / len(idx)
    w -= 0.1 * g

Practical Examples

1. Linear regression by hand vs solver

# analytic
w_closed = np.linalg.lstsq(Xb, y, rcond=None)[0]
# gradient descent (10 lines) converges to the same answer
w = np.zeros(Xb.shape[1])
for i in range(2000):
    w -= 0.01 * Xb.T @ (Xb @ w - y) / len(y)
np.allclose(w, w_closed, atol=1e-2)

2. Effect of learning rate

for lr in [0.5, 0.1, 0.01, 0.001]:
    w = 0.0
    traj = []
    for _ in range(50):
        w -= lr * (2 * (w - 3))          # minimize (w-3)^2
        traj.append(w)
    print(lr, "final:", traj[-1], "oscillates:", max(traj) > 6)
    # 0.5 overshoots, 0.1 fine, 0.001 too slow for 50 steps

3. Full training loop with AdamW + cosine schedule

import torch, torch.nn as nn
from torch.optim.lr_scheduler import CosineAnnealingLR

model = nn.Linear(10, 1)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
sched = CosineAnnealingLR(opt, T_max=100)

for epoch in range(100):
    opt.zero_grad()
    loss = nn.functional.mse_loss(model(X), y)
    loss.backward()
    opt.step()
    sched.step()

Common Patterns

# Gradient clipping (RNNs/transformers)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# Feature scaling before GD (conditioning)
X = (X - X.mean(0)) / X.std(0)        # poor scaling -> zigzag convergence

# Watch the loss curve, not just the final number
# plot train & val loss per epoch -> diagnose lr, overfit, plateau

# Early stopping as implicit regularization
if val_loss < best_val: best_val = val_loss; torch.save(model.state_dict(), "best.pt")

Best Practices

  • Start simple: plain SGD or Adam with defaults before exotic schedules.
  • Track training + validation loss curves — flat train + rising val = overfit; both flat = lr too low.
  • Scale features (mean 0, unit variance) so gradient steps move all directions equally.
  • Check gradient correctness once (finite differences) when writing custom losses.
  • Warm up + cosine decay for transformers; constant small lr for quick experiments.
  • Weight decay ≠ L2 in Adam: use AdamW; it decouples decay from the adaptive step.
  • Random/Bayesian beats grid search beyond ~3 hyperparameters.

Common Mistakes

Mistake Consequence Fix
Learning rate too high Loss diverges to NaN Lower $\eta$; log-scale search
Learning rate too low Looks stuck, slow progress Raise or use schedule; check curve
No feature scaling Slow zigzag convergence Standardize inputs
Treating neural loss as convex Confusion at local minima Accept local minima; tune init/seed
Applying L2 to Adam Wrong decay behavior AdamW with weight_decay
Ignoring res.success from minimize Trusting a failed optimum Check return flags
Grid search with 5+ params Combinatorial explosion Random/Optuna

Troubleshooting

Symptom Cause Fix
Loss = NaN early lr too high, bad init, log(0) Lower lr; check inputs; add epsilon
Loss flatlines forever lr too low / saturated layer Raise lr or change activation
Val loss worse than train consistently Overfitting Regularize, early stop, more data
Converges but slowly Conditioning (unequal scales) Standardize; use Adam/momentum
Diverges only after epoch 30 Scheduler decay too aggressive Gentler schedule; check curve

Official Documentation

Last Updated

2026-09-06 — Initial version.