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 are navigating a complex loss surface, so 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, so 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#

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, so no global optimality guarantees apply
• 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 surface, 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#

What next#

• Advanced optimization. Second-order methods (Newton, L-BFGS), variance reduction (SVRG), distributed training, sharpness-aware minimisation (SAM), and automatic differentiation, for readers who want Boyd-level depth and ML-frontier training science.
• Probability and 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.

Knowledge check#

Browser lab: a training-diagnostics dashboard#

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))

# 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}")
print("Notice: both dashboards answer the same four questions — convexity, smoothness, conditioning, and noise.")

Key takeaways#

• Method choice follows from five questions: convex or not, smoothness $L$, condition number $\kappa$, gradient noise, and constraint structure.
• Each answer maps to a concrete optimizer, step size, and batch size; the method-selection table is this course compressed into one page.
• The 2026 default is AdamW with cosine decay and warmup; switch to SGD + momentum + cosine when chasing final accuracy on a well-understood task.
• RL policy gradients are nonconvex, high-variance, and non-stationary, so they use Adam with a small $\eta$ and gradient clipping.
• Diffusion training is a large but benign least-squares-style problem, so it uses AdamW with cosine decay plus warmup and a weight EMA.
• Every earlier week shows up in the checklist: convexity (W2), smoothness (W3), conditioning (W4), stochasticity (W6), momentum and adaptivity (W7-8), schedules and clipping (W9), proximal and constraints (W10-11), duality (W12), and nonconvex practice (W13).

Further reading#

• Ian Goodfellow, Yoshua Bengio & Aaron Courville, Deep Learning (MIT Press, 2016). Chapter 8 is the practitioner's guide to optimiser selection and training diagnostics.
• Léon Bottou, Frank E. Curtis & Jorge Nocedal, Optimization Methods for Large-Scale Machine Learning (SIAM Review, 2018). Sections 6–7 discuss method choice and diagnostics at scale.
• Diederik P. Kingma & Jimmy Ba, "Adam: A Method for Stochastic Optimization" (ICLR 2015). The default-optimiser reference. Freely available at arXiv:1412.6980.
• Ilya Loshchilov & Frank Hutter, "Decoupled Weight Decay Regularization" (ICLR 2019). AdamW, the modern default. Freely available at arXiv:1711.05101.