Calculus and optimisation

Everything that “trains” is minimising a loss by walking downhill. You need enough calculus to explain why training diverges, why it plateaus, and what the optimiser knobs actually do.

Gradient: the direction of steepest ascent

For f(w) with w a vector, the gradient ∇f is the vector of partial derivatives. It points uphill, so we step against it.

w = w - lr * grad          # gradient descent, the whole idea

The learning rate lr is the step size, and it is the hyperparameter that most often decides whether training works at all:

NaN loss a few steps into training is almost always learning rate, exploding gradients, or a log(0) in your loss. Check in that order.

Chain rule = backpropagation

For f(g(x)), df/dx = df/dg * dg/dx. Stack that through every layer and you have backprop: compute the loss, then propagate derivatives backwards, multiplying local gradients.

# Forward
z1 = X @ W1 + b1
a1 = relu(z1)
z2 = a1 @ W2 + b2
loss = mse(z2, y)

# Backward - each step is the chain rule applied locally
dz2 = 2 * (z2 - y) / len(y)
dW2 = a1.T @ dz2
da1 = dz2 @ W2.T
dz1 = da1 * (z1 > 0)          # ReLU derivative: 1 where positive, else 0
dW1 = X.T @ dz1

Why gradients vanish. Each layer multiplies by a local derivative. If those are consistently < 1, the product shrinks exponentially with depth and early layers stop learning. Sigmoid saturates at derivative ~0.25 at best, which is why deep sigmoid networks were untrainable and why ReLU (derivative exactly 1 on the positive side) unlocked depth. Residual connections attack the same problem by giving the gradient an identity path.

Why they explode. The mirror case — local derivatives > 1 compound. Fix with gradient clipping:

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

Standard practice in RNN and transformer training.

Convexity, and why it mostly doesn’t matter any more

A convex function has exactly one minimum, so gradient descent provably finds it. Linear and logistic regression with L2 are convex — that’s why they’re reliable and reproducible.

Neural networks are wildly non-convex, with countless local minima and saddle points. In practice this matters far less than the theory suggests: in high dimensions, most critical points are saddles rather than bad local minima, and the many minima that exist tend to have similar loss. The practical consequence is that training is stochastic — two runs with different seeds give different weights and similar performance.

The optimiser family

The AdamW distinction is a real interview question. In Adam, L2 regularisation gets folded into the gradient and then divided by the adaptive scaling, which weakens it inconsistently across parameters. AdamW applies decay directly to the weights instead, so it behaves as intended. Every modern transformer recipe uses AdamW.

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)

Memory cost worth knowing: Adam stores two extra tensors per parameter (first and second moment). That’s roughly 3x the model size just for optimiser state, which is why full fine-tuning of a large model needs so much more memory than inference, and why LoRA and 8-bit optimisers exist.

Batch size and the noise trade-off

Rule of thumb when scaling up: increase learning rate roughly with batch size (linear scaling), and add warmup so the first steps don’t blow up.

Learning-rate schedules

Constant learning rates are rarely optimal. The standard transformer recipe is linear warmup then cosine decay:

from torch.optim.lr_scheduler import OneCycleLR

scheduler = OneCycleLR(optimizer, max_lr=3e-4,
                       total_steps=total_steps, pct_start=0.05)

Warmup exists because early gradients are large and Adam’s second-moment estimate is still unreliable; stepping at full rate immediately destabilises training. Decay exists because you want big steps to find the basin and small steps to settle into it.

Loss functions and what they encode

Choosing MSE when the target has outliers is a common self-inflicted wound: a handful of extreme labels dominate the gradient and the model fits them at everyone else’s expense. Huber or log-transforming the target usually fixes it.

Cross-entropy connects directly to information theory — see 04_information_theory.md.

Interview angle

• “Explain backpropagation.” — apply the chain rule layer by layer from the loss backwards, reusing each layer’s local derivative. It’s not a learning algorithm, it’s an efficient way to compute gradients; gradient descent is what learns.
• “Your loss becomes NaN after 50 steps. Debug it.” — learning rate too high, exploding gradients, or a numerical hazard (log(0), division by a near-zero, fp16 overflow). Lower the LR by 10x, add gradient clipping, check the loss for domain errors, and confirm inputs are normalised.
• “Why did deep networks with sigmoid fail to train?” — vanishing gradients. Sigmoid’s derivative peaks at 0.25 and saturates near 0 at both ends, so the product across many layers collapses. ReLU and residual connections give the gradient a path with derivative 1.
• “Adam vs AdamW?” — AdamW decouples weight decay from the adaptive gradient scaling. In plain Adam, L2 gets divided by the per-parameter scale and no longer regularises uniformly. AdamW is the transformer default.
• “Why is full fine-tuning so much more memory-hungry than inference?” — weights, plus gradients, plus Adam’s two moment tensors per parameter, plus activations kept for the backward pass. Roughly 4x the parameter memory before activations. It’s the motivation for LoRA, gradient checkpointing and 8-bit optimisers.
• “Why warmup?” — early gradients are large and Adam’s variance estimate is not yet calibrated; stepping at full LR immediately can destabilise or diverge training.
• “Small vs large batches?” — small batches give noisy gradients that often generalise better and use hardware poorly; large batches are efficient and stable but may need LR scaling and warmup to match generalisation.