Week 14: Capstone — Theory Meets Training

The Full Diagnostic Checklist#

Before choosing an optimiser, run through this checklist:

1. Convex or not? If the loss is convex (linear/logistic regression, SVM dual), global optimality is achievable. If nonconvex (deep nets), you're navigating a landscape — focus on architecture and initialisation.

2. Smoothness estimate: What is $L$, the Lipschitz constant of the gradient? For quadratics $f(x) = \frac{1}{2}x^T Q x$, $L = \lambda_{\max}(Q)$. For deep nets, estimate via a power-iteration sweep on a mini-batch Hessian. Start with $\eta \approx 1/L$.

3. Condition number: Estimate $\kappa = \lambda_{\max} / \lambda_{\min}$. Large $\kappa$ ($>100$) means GD will be slow — add momentum or use Adam.

4. Noise level: What is the variance of mini-batch gradients? High noise (small $B$, heterogeneous data) → use momentum or Adam + warmup + gradient clipping. Low noise (large $B$, homogeneous data) → plain SGD may suffice.

5. Constraint structure: Are there box constraints, a simplex constraint, or a norm bound? If projections are cheap (ball, box), use projected GD. If not, consider soft penalties.

Method Selection Table#

| Problem profile | Recommended optimiser | Step size | Batch size $B$ | |---|---|---|---| | Small, convex, well-conditioned ($\kappa < 10$) | GD | $\eta = 1/L$ | $n$ (full) | | Large $n$, convex, moderate $\kappa$ | SGD | $\eta_0 = 1/L$, decay $\propto 1/\sqrt{k}$ | 32–256 | | Large $d$, convex, ill-conditioned ($\kappa > 100$) | SGD + momentum | $\eta \approx 1/L$, $\beta = 0.9$ | 32–256 | | Large $n$, nonconvex, sparse gradients | Adam | $\eta = 10^{-3}$, defaults | 32–512 | | Nonconvex, best generalisation needed | SGD + momentum + cosine LR | tuned $\eta$, $\beta = 0.9$ | 256–1024 | | Constrained (ball/box/simplex) | Projected GD/SGD | $\eta = 1/L$ for GD | as above | | Composite (smooth + $\ell_1$) | Proximal gradient (ISTA) | $\eta = 1/L_g$ | as above |

The default in 2026: Adam/AdamW with cosine LR schedule and small warmup. This is the safe choice when prototyping. Switch to SGD + momentum + cosine only when maximising final performance on a well-understood problem.

Case Study A: Policy Gradient for RL#

Reinforcement learning optimises a policy $\pi_\theta(a \mid s)$ to maximise expected return. The policy gradient estimator is:

$\nabla_\theta J(\theta) \approx \frac{1}{B} \sum_{i=1}^B \nabla_\theta \log \pi_\theta(a_i \mid s_i) \cdot \hat A_i$

where $\hat A_i$ is an advantage estimate.

Optimisation profile:

• Nonconvex: The policy objective is highly nonconvex — no global optimality guarantees
• High variance: The advantage estimates $\hat A_i$ are noisy, creating large gradient variance even with $B$ moderate
• Non-stationary: The data distribution shifts as the policy changes, and value function estimates co-evolve

Standard recipe: Adam with $\eta = 3 \times 10^{-4}$, $\beta_1 = 0.9$, $\beta_2 = 0.999$, batch size determined by environment parallelism. Gradient clipping at norm $0.5$–$1.0$. Linear LR decay or constant $\eta$. The value function baseline acts as a variance-reduction mechanism.

Link: Reinforcement Learning course

Case Study B: Denoising Diffusion Training#

Diffusion models and flow matching train a denoiser $\epsilon_\theta(x_t, t)$ to predict the noise added to a clean sample. The training objective is:

$J(\theta) = \mathbb{E}_{t, x_0, \epsilon} \left[ \|\epsilon - \epsilon_\theta(\sqrt{\bar\alpha_t}\, x_0 + \sqrt{1-\bar\alpha_t}\,\epsilon, t)\|_2^2 \right]$

Optimisation profile:

• Large-scale: Models are large ($10^7$–$10^9$ parameters), datasets enormous
• Well-behaved objective: The L2 denoising loss has a relatively benign landscape — close to a massive least-squares problem in function space
• Low gradient noise: The per-example loss is a simple squared error with explicit targets

Standard recipe: AdamW (Adam with decoupled weight decay) with $\eta = 10^{-4}$, $\beta_1 = 0.9$, $\beta_2 = 0.999$, cosine LR schedule with warmup. EMA of model weights (not gradient moments) for final inference. Large batch $B \approx 256$–$2048$ distributed across GPUs.

Link: Generative Models course

Course Map#

| Concept | Where | Key equation | |---|---|---| | ERM | W1 | $J(\theta) = \frac{1}{n}\sum \ell(f_\theta(x_i), y_i)$ | | Convexity | W2 | $f(\theta x + (1-\theta)y) \leq \theta f(x) + (1-\theta)f(y)$ | | Smoothness & rates | W3 | $f(x_K) - f^* \leq \frac{2L\|x_0-x^*\|^2}{K}$ | | Least squares | W4 | $\min_w \|Xw-y\|_2^2$, $\kappa = \lambda_{\max}/\lambda_{\min}$ | | Logistic / softmax | W5 | $\nabla_w \ell = (\sigma(w^T x) - y) x$ | | SGD | W6 | $w_{k+1} = w_k - \eta_k g_k$ (mini-batch) | | Momentum / Nesterov | W7 | $v_{k+1} = \beta v_k - \eta \nabla f(w_k)$ | | Adam | W8 | $m_k = \beta_1 m_{k-1} + (1-\beta_1) g_k$, $v_k = \beta_2 v_{k-1} + (1-\beta_2) g_k^2$ | | LR schedules / clipping | W9 | Cosine, warmup, $\|g\| \leq C$ | | Subgradients / proximal | W10 | $x_{k+1} = \text{prox}_{\eta h}(x_k - \eta \nabla g(x_k))$ | | Constraints / projections | W11 | $x_{k+1} = P_C(x_k - \eta \nabla f(x_k))$ | | Duality | W12 | $g(\lambda) = \min_x L(x, \lambda)$, SVM dual | | Nonconvex | W13 | Overparam, init, normalisation |

What Next#

• Advanced Optimization: Second-order methods (Newton, L-BFGS), variance reduction (SVRG), distributed training, sharpness-aware minimisation (SAM), and automatic differentiation — for those who want Boyd-level depth and ML-frontier training science.
• Probability & Statistics: The modelling side — MLE/MAP, exponential families, graphical models, and information theory (entropy/KL) — the language connecting loss functions to probabilistic stories.
• Reinforcement Learning and Generative Models: These applied courses use the optimisers from this course as black-box subroutine drivers. You now understand what's inside the black box.

Browser lab#

A mini training-diagnostics dashboard: (1) estimate $L$ (smoothness) on a synthetic LS problem, plot loss + gradient norm; (2) run Adam on a small nonconvex problem and print the diagnostic checklist output.

import numpy as np
import matplotlib.pyplot as plt

rng = np.random.default_rng(42)

# --- Part 1: Estimate L on least squares ---
n, d = 200, 5
X = rng.normal(size=(n, d))
w_true = rng.normal(size=d)
y = X @ w_true + rng.normal(scale=0.2, size=n)

L_est = np.linalg.eigvalsh(X.T @ X / n).max()
print(f"[Diagnostic] Smoothness estimate: L ≈ {L_est:.4f}")
print(f"[Diagnostic] Recommended η for GD: 1/L ≈ {1/L_est:.4f}")

w = np.zeros(d)
eta = 1.0 / L_est
losses, grad_norms = [], []
for _ in range(60):
    resid = X @ w - y
    g = X.T @ resid / n
    w = w - eta * g
    losses.append(0.5 * np.mean(resid**2))
    grad_norms.append(np.linalg.norm(g))

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4.5))
ax1.semilogy(losses); ax1.set_title("Loss"); ax1.set_xlabel("Iter"); ax1.grid(True, alpha=0.3)
ax2.semilogy(grad_norms); ax2.set_title("||grad||"); ax2.set_xlabel("Iter"); ax2.grid(True, alpha=0.3)

# --- Part 2: Nonconvex check with Adam ---
def f_nonconvex(w):
    x, y = w[0], w[1]
    return x**4 - 3*x**2 + 0.5*x + y**4 - 2*y**2 + 0.3*y

def grad_nonconvex(w):
    x, y = w[0], w[1]
    return np.array([4*x**3 - 6*x + 0.5, 4*y**3 - 4*y + 0.3])

print("\n[Diagnostic] Nonconvex problem: f(x,y) = x⁴−3x²+½x + y⁴−2y²+0.3y")
print("[Diagnostic] Status: NONCONVEX — multiple local minima")
print("[Diagnostic] Method: Adam (default), η=0.05")
print("[Diagnostic] No constraints, moderate noise expected")

w = np.array([2.0, 2.0])
m = np.zeros(2); v = np.zeros(2)
eta_a, b1, b2, eps = 0.05, 0.9, 0.999, 1e-8
losses_nc = []
for t in range(1, 151):
    g = grad_nonconvex(w)
    m = b1*m + (1-b1)*g; v = b2*v + (1-b2)*g**2
    mh = m/(1-b1**t); vh = v/(1-b2**t)
    w = w - eta_a * mh/(np.sqrt(vh)+eps)
    losses_nc.append(f_nonconvex(w))

ax_nc = plt.gcf().add_subplot(1, 2, 2)
# Overwrite ax2 with the nonconvex results
ax2.clear()
ax2.plot(losses_nc)
ax2.set_title("Nonconvex Loss (Adam)"); ax2.set_xlabel("Iter")
ax2.grid(True, alpha=0.3)

plt.tight_layout(); plt.show()
print(f"\n[Diagnostic] Final loss: {losses_nc[-1]:.6f}")