Non-convex optimisation

Convex optimisation would be kinder. Language-model loss is not convex. You will still use first-order methods because the replica that matters is the one you can run, not the one you can solve in closed form.

Topics

• Gradient descent, momentum, and Adam. Write the update, then say what each buffer stores.
• Learning-rate schedules. Warmup exists because early gradients are noisy, not because a blog said so.
• Loss sharpness and the Hessian. You will not form the full Hessian. You can still talk about curvature.
• Exploding and vanishing signals. Clip, residual paths, and careful init are engineering answers to geometry.
• Generalisation folklore versus what you will log: train loss, grad norm, tokens per second, and a held-out NLL.

Adam keeps a biased first moment and a biased second moment of the gradient. The epsilon in the denominator is not decoration. When a second-moment entry is tiny, the effective step is huge. That is one real source of training instability.

Coding

import math
import torch

def adam_step(p, g, m, v, t, lr=3e-4, beta1=0.9, beta2=0.95, eps=1e-8, wd=0.1):
    m = beta1 * m + (1 - beta1) * g
    v = beta2 * v + (1 - beta2) * (g * g)
    m_hat = m / (1 - beta1 ** t)
    v_hat = v / (1 - beta2 ** t)
    p = p * (1 - lr * wd) - lr * m_hat / (torch.sqrt(v_hat) + eps)
    return p, m, v

def cosine_lr(step, total, base, warmup):
    if step < warmup:
        return base * step / max(1, warmup)
    progress = (step - warmup) / max(1, total - warmup)
    return base * 0.5 * (1.0 + math.cos(math.pi * progress))

Run a toy non-convex scalar, f(x) = x^4 - x^2, from several starts. Plot the path. Then run a two-layer MLP on a noisy spiral and log grad norms. You want to see a spike and a clip, on purpose.

Assignment

• Implement SGD, momentum, and Adam on the same toy loss. Do not claim a winner from one seed. Report three seeds.
• Break Adam: set eps to 0 on a first step with a near-zero gradient and write what happens.
• Add global grad clip at 1.0 and show the clip-to-unclipped ratio for 200 steps.

Opinion checkpoint

Papers you may name by title when you write notes: none required this week. You will meet training instability again when you pretrain in week 5.

Next: Build a transformer from scratch (TinyGPT).