Training a network: losses, mini-batches and the training loop

Before this lesson: lesson 01 (the MLP); logistic and softmax regression and the train/validation/test split (ml-fundamentals track); gradient descent (optimization track); maximum likelihood (probability-and-statistics track); cross-entropy (information-theory track). Eigenvalues appear once (linear algebra, lesson 05). Time: about 65 minutes, plus the notebook.

---

Why this matters #

Every model in this track is trained with the same loop: forward pass, loss, gradients, optimizer step. ImageNet classifiers, GPT-style language models (next-token cross-entropy over a vocabulary of about 50,000 tokens), diffusion models (a mean squared error) and physics-informed networks (a sum of squared residuals) differ in the model and the loss, not in the loop. Most failed training runs come from a few mistakes in this lesson: probabilities passed where logits are expected, a missing zero_grad(), a learning rate above the stability limit, or model.eval() forgotten at test time. The next lesson (backpropagation) explains how the gradients are computed; this one explains what is done with them.

1. Predict first #

2. The picture #

Training is a loop #

A training run repeats four steps:

1. take a small random mini-batch of training examples;
2. run the network on it (the forward pass) and compute the average loss;
3. compute the gradient of that loss with respect to every weight (lesson 03 shows how);
4. move every weight a small step against its gradient: $\theta \leftarrow \theta - \eta\, g$.

Every few steps we also compute the loss on validation data that the network never trains on. The training loss tells us whether the network can fit the data. The validation loss tells us whether it works on new data.

Interactive figure: loading…

Each step is noisy, but right on average #

The true direction of descent is the full gradient: the average of the gradients of all $n$ training examples. A mini-batch of $B$ examples gives an average over only $B$ of them. Different mini-batches give different directions, scattered around the full gradient. Larger batches scatter less, but each step costs more.

Interactive figure: loading…

The step size has a limit #

The simplest loss is a parabola, $L(w) = \frac{\lambda}{2}w^2$, with curvature $\lambda$. Its gradient is $\lambda w$, so one gradient descent step gives $w \leftarrow w - \eta\lambda w = (1 - \eta\lambda)w$. Every step multiplies $w$ by the same number $1 - \eta\lambda$. If this number is between $-1$ and $1$, $w$ shrinks to 0. If it is below $-1$, $w$ changes sign and grows at every step: the loss explodes.

Interactive figure: loading…

3. Definitions and main results #

3.1 The learning problem #

A network $f_\theta$ has parameters $\theta$ (all its weights and biases). The data are pairs $(x_i, y_i)$, $i = 1, \dots, n$: an input and its target. A loss function $\ell(\hat y, y)$ measures how wrong a prediction $\hat y$ is for the target $y$. Training means minimizing the empirical risk, the average loss over the training set:

What we really care about is the population risk $R(\theta) = \mathbb{E}_{(x, y)}\,\ell(f_\theta(x), y)$: the expected loss on new examples from the same distribution. We cannot compute $R$, so we estimate it with held-out data. The data are split into three sets with three different jobs:

Set	Used for	Rule
training set	computing gradients and updating $\theta$	the only data the optimizer sees
validation set	choosing hyperparameters (width, learning rate, number of steps) and early stopping	never used for gradients
test set	one final estimate of $R$	used once, at the end; choices based on it make it optimistic

Three different "losses". The mini-batch loss is the average over the $B$ examples of one step; it is noisy. The training loss $\hat R(\theta)$ is the average over the whole training set. The validation (or test) loss is the same average on held-out data. Do not mix them up when you read a loss curve.

3.2 Output layers and losses from maximum likelihood #

Let the network output the parameters of a probability distribution $p_\theta(y \mid x)$ for the target. Maximum likelihood chooses $\theta$ to make the observed targets as probable as possible. For independent examples this means maximizing $\prod_i p_\theta(y_i \mid x_i)$, which is the same as minimizing the negative log-likelihood $-\frac1n \sum_i \log p_\theta(y_i \mid x_i)$. Each loss below is this recipe for one choice of distribution.

Regression: Gaussian noise gives the squared error. Assume $y = f_\theta(x) + \text{noise}$, with Gaussian noise of variance $\sigma^2$. Then

The second term does not depend on $\theta$. So minimizing the negative log-likelihood is minimizing the mean squared error (MSE). Output layer: linear, no activation.

Binary classification: Bernoulli gives binary cross-entropy. The network outputs one logit $z \in \mathbb{R}$, and $s(z) = 1/(1 + e^{-z})$ (the logistic sigmoid) is the probability of class 1. For $y \in \{0, 1\}$, $p(y) = s^y (1 - s)^{1 - y}$, so

the binary cross-entropy (BCE). Its gradient is short. Use $\frac{d}{dz}\log s = 1 - s$ and $\frac{d}{dz}\log(1 - s) = -s$:

$K$ classes: softmax and cross-entropy. The network outputs $K$ logits $z \in \mathbb{R}^K$. The softmax turns them into probabilities:

With a one-hot target $y$ (1 at the true class, 0 elsewhere) the negative log-likelihood is the cross-entropy

where $c^*$ is the true class. It is the cross-entropy between the target distribution $y$ and the predicted distribution $p$ (information-theory track).

Theorem (gradient of softmax cross-entropy). For any target $y$ with $\sum_c y_c = 1$,

Proof.

1. $\log p_c = z_c - \log \sum_j e^{z_j}$, and $\frac{\partial}{\partial z_k} \log \sum_j e^{z_j} = \frac{e^{z_k}}{\sum_j e^{z_j}} = p_k$.
2. So $\frac{\partial \ell}{\partial z_k} = -\sum_c y_c\,(\delta_{ck} - p_k) = -y_k + p_k \sum_c y_c = p_k - y_k$. $\blacksquare$

The gradient is "predicted minus target". Its entries sum to $\sum_k p_k - \sum_k y_k = 0$. For a mini-batch the loss is the mean over the $B$ rows, so the gradient with respect to the logit matrix is $(P - Y)/B$; lesson 03 starts the backward pass from there. The binary case is the same result with $K = 2$.

Why not a sigmoid followed by the squared error? Take one logit, target $y = 1$, and the loss $\ell = \frac12 (s(z) - y)^2$. The chain rule with $s'(z) = s(1 - s)$ gives

When the model is confidently wrong ($z \ll 0$, so $s \approx 0$), the factor $s(1 - s)$ is close to 0, and the gradient almost vanishes. At $z = -5$: $s = 0.0067$, the BCE gradient is $s - 1 = -0.993$, and the MSE gradient is $-0.0066$, about 150 times smaller. The ratio is always $1/(s(1 - s)) \ge 4$. Cross-entropy learns fastest from its worst mistakes; sigmoid + MSE learns slowest from them.

Interactive figure: loading…

3.3 Numerically stable losses #

Computing the softmax as written fails for large logits. In float64, $e^{z}$ overflows to $\infty$ for $z > 709.8$; in float32 already for $z > 88.7$. For very negative logits, $e^{z}$ underflows to 0, and $\log 0 = -\infty$.

Softmax ignores a common shift. For any number $a$, $\operatorname{softmax}(z + a\mathbf{1}) = \operatorname{softmax}(z)$, because $e^{z_c + a} / \sum_j e^{z_j + a} = e^a e^{z_c} / (e^a \sum_j e^{z_j})$. So we may subtract the largest logit $m = \max_j z_j$ before exponentiating. This is the log-sum-exp trick:

Now every exponent is $\le 0$, so nothing overflows, and the largest term is $e^0 = 1$, so the sum is at least 1 and its log is finite. The loss is computed as $\log p$ directly (log-softmax), never as the log of a probability that may have underflowed.

This is why PyTorch losses take logits. nn.CrossEntropyLoss expects raw logits and applies log-softmax inside; nn.BCEWithLogitsLoss expects the logit $z$ and uses the stable form $\max(z, 0) - zy + \log(1 + e^{-|z|})$. Example in float32 with logits $(0, -110)$: the softmax gives $p_2 = 0$ (underflow) and $\log p_2 = -\infty$, while log-softmax gives exactly $-110$.

3.4 Gradient descent and mini-batch SGD #

Gradient descent (GD) updates $\theta \leftarrow \theta - \eta \nabla \hat R(\theta)$ with a learning rate $\eta > 0$. One step needs the gradient of all $n$ examples: too expensive when $n$ is in the millions.

Mini-batch stochastic gradient descent (SGD) picks a random set $\mathcal{B}$ of $B$ indices and uses

Theorem (the mini-batch gradient is unbiased, with variance $\Sigma/B$). Let the $B$ indices be drawn independently and uniformly from $\{1, \dots, n\}$, and let $\Sigma = \frac1n \sum_i (\nabla \ell_i - \nabla \hat R)(\nabla \ell_i - \nabla \hat R)^\top$ be the covariance of the per-example gradients. Then

Proof.

1. One uniformly drawn index $i$ gives $\mathbb{E}[\nabla \ell_i] = \frac1n \sum_{j} \nabla \ell_j = \nabla \hat R$, and by the definition of $\Sigma$, $\operatorname{Cov}(\nabla \ell_i) = \Sigma$.
2. The mean of the $B$ terms has the same expectation: $\mathbb{E}[g_{\mathcal{B}}] = \nabla \hat R$.
3. The $B$ terms are independent, so the covariance of their sum is the sum of the covariances, $B\Sigma$. Dividing the sum by $B$ divides the covariance by $B^2$: $\operatorname{Cov}(g_{\mathcal{B}}) = B\Sigma/B^2 = \Sigma/B$. $\blacksquare$

So the standard deviation of each gradient entry falls like $1/\sqrt{B}$. In practice we shuffle the data and take batches without replacement; the mini-batch gradient is still unbiased, and the variance is a little smaller: $\frac{\Sigma}{B} \cdot \frac{n - B}{n - 1}$.

Vocabulary. One iteration (or step) uses one mini-batch. One epoch is one pass over the whole training set: $n/B$ iterations. Shuffling the order every epoch keeps the batches random.

Why noise is acceptable. Each step is $n/B$ times cheaper than a full-batch step, and it points in the right direction on average. Early in training the full gradient is large compared with the noise, so noisy steps make fast progress. Near a minimum the full gradient becomes small and the noise dominates (you saw this in the figure); the iterates then jump around the minimum, which is one reason to lower the learning rate late in training (lesson 12).

Adam in one paragraph. Most networks in this track are trained with Adam (Kingma and Ba 2015). It keeps two running averages per parameter: $m \leftarrow \beta_1 m + (1 - \beta_1) g$ (a smoothed gradient, momentum) and $v \leftarrow \beta_2 v + (1 - \beta_2) g^2$ (a smoothed squared gradient). The update is $\theta \leftarrow \theta - \eta\, \hat m / (\sqrt{\hat v} + \epsilon)$, where $\hat m, \hat v$ are the averages divided by $1 - \beta_1^t$ and $1 - \beta_2^t$ (bias correction at step $t$). Dividing by $\sqrt{\hat v}$ gives every parameter a step of similar size, whatever the scale of its gradient. The defaults are $\beta_1 = 0.9$, $\beta_2 = 0.999$, $\epsilon = 10^{-8}$; a common learning rate is $10^{-3}$. The optimization track derives SGD with momentum, Adam and AdamW in detail.

3.5 The learning rate on a quadratic #

Theorem. Gradient descent on $L(w) = \frac{\lambda}{2} w^2$ with $\lambda > 0$ converges to 0 from every $w_0 \ne 0$ if and only if $0 < \eta < 2/\lambda$.

Proof.

1. $L'(w) = \lambda w$, so $w_{t+1} = w_t - \eta \lambda w_t = (1 - \eta\lambda)\, w_t$, and $w_t = (1 - \eta\lambda)^t w_0$.
2. $(1 - \eta\lambda)^t \to 0$ if and only if $|1 - \eta\lambda| < 1$, that is, $-1 < 1 - \eta\lambda < 1$.
3. The right inequality gives $\eta > 0$; the left one gives $\eta < 2/\lambda$. $\blacksquare$

The fastest choice is $\eta = 1/\lambda$ (one step). For $1/\lambda < \eta < 2/\lambda$ the iterates change sign every step (a zig-zag) but still converge. For a quadratic in many dimensions with Hessian $H$, the same argument runs along each eigenvector of $H$ (linear algebra, lesson 05), so the limit is $\eta < 2/\lambda_{\max}(H)$: the most curved direction sets the limit. For a network, $\lambda_{\max}$ of the Hessian (the sharpness) changes during training; lesson 12 shows how warmup and learning-rate schedules deal with this.

3.6 The PyTorch training loop #

Here is the complete loop for a classifier. Each line is explained below.

model = nn.Sequential(nn.Linear(2, 32), nn.ReLU(), nn.Linear(32, 2))  # outputs logits
loss_fn = nn.CrossEntropyLoss()                    # log-softmax + NLL, takes logits
opt = torch.optim.Adam(model.parameters(), lr=1e-2)

for epoch in range(30):
    model.train()                                  # training behaviour (dropout on, BatchNorm batch stats)
    for xb, yb in train_loader:                    # shuffled mini-batches
        logits = model(xb)                         # 1. forward pass
        loss = loss_fn(logits, yb)                 # 2. mean loss over the batch
        opt.zero_grad()                            # 3. clear the old gradients
        loss.backward()                            # 4. backprop: fills p.grad for every parameter
        opt.step()                                 # 5. update the parameters
    model.eval()                                   # evaluation behaviour (dropout off, BatchNorm running stats)
    with torch.no_grad():                          # no graph, no gradients: faster, less memory
        val_logits = model(X_val)
        val_loss = loss_fn(val_logits, y_val).item()
        val_acc = (val_logits.argmax(dim=1) == y_val).float().mean().item()

• opt.zero_grad(). PyTorch adds new gradients to p.grad; it does not overwrite them. This is useful for gradient accumulation (several small batches before one step), but in a normal loop the old gradient must be cleared. Without it, step $t$ uses the sum of all gradients so far, and training usually becomes unstable.
• loss.backward() runs backpropagation (lesson 03) from the scalar loss.
• opt.step() applies the update rule (SGD, Adam, …) to every parameter, using p.grad.
• model.train() / model.eval() switch the behaviour of layers such as dropout (lesson 06) and BatchNorm (lesson 05). They do not turn gradients on or off.
• torch.no_grad() turns off the recording of the computation graph. Use it for evaluation; it does not change dropout or BatchNorm.

3.7 Reading loss curves #

What you see	Likely cause	What to do
training and validation loss both high and flat	underfitting: model too small, too few steps, or learning rate too small	bigger model, more steps, larger $\eta$
training loss falls, validation loss falls, then rises	overfitting: the model fits noise in the training set	stop early, more data, regularization (lesson 06)
loss jumps up and down, rises above its start, or becomes NaN	learning rate too large	lower $\eta$ (a factor 3 to 10), warmup, gradient clipping (lesson 12)
loss stays at $\log K$	the model outputs uniform probabilities: a bug (labels shuffled, no gradient reaches the weights, learning rate 0)	check the data and that p.grad is nonzero
loss falls but stays above $\log(1 + (K - 1)/e)$	softmax applied before CrossEntropyLoss	pass logits

Sanity checks before a long run.

1. Initial loss. With near-zero logits, $p = (1/K, \dots, 1/K)$ and the loss is $\log K$: $0.693$ for 2 classes, $2.303$ for 10, $6.91$ for 1000. A very different first loss points to a bug or a bad initialization.
2. Overfit one small batch. Train on a single batch of a few examples. A correct model and loop drive that loss close to 0. If it cannot, there is a bug.
3. Watch both curves. A training loss alone cannot show overfitting.

Answers to section 1. (1) About $\log 10 \approx 2.30$: logits near 0 give $p \approx 1/10$ for every class, and the loss is $-\log(1/10)$. (2) It halves: the variance falls by $128/32 = 4$, the standard deviation by $\sqrt 4 = 2$. (3) $L(w) = w^2$ has $\lambda = 2$ and limit $2/\lambda = 1$. With $\eta = 1.1$, $w_{t+1} = (1 - 2.2)\,w_t = -1.2\,w_t$: $w$ changes sign and grows by 20% at every step. It diverges.

4. Worked example #

(a) One softmax cross-entropy step by hand. Three classes, logits $z = (2, 1, -1)$, true class 1, so $y = (1, 0, 0)$.

1. Exponentials: $e^2 = 7.389$, $e^1 = 2.718$, $e^{-1} = 0.368$. Their sum is $10.475$.
2. Probabilities: $p = (7.389, 2.718, 0.368)/10.475 = (0.7054, 0.2595, 0.0351)$.
3. Loss: $\ell = -\log 0.7054 = 0.3490$.
4. Gradient: $p - y = (0.7054 - 1,\ 0.2595,\ 0.0351) = (-0.2946, 0.2595, 0.0351)$. The entries sum to 0.
5. Suppose we could move the logits directly, with $\eta = 1$: $z \leftarrow z - (p - y) = (2.2946, 0.7405, -1.0351)$. The new probabilities are $(0.8018, 0.1695, 0.0287)$ and the new loss is $-\log 0.8018 = 0.2209$. The step raised the true logit, lowered the others, and the loss fell from $0.349$ to $0.221$.

(b) The same loss, stably. Add 998 to every logit: $z = (1000, 999, 997)$. A naive softmax computes $e^{1000} = \infty$ and fails. With $m = 1000$:

and $\ell = 1000.3490 - z_1 = 1000.3490 - 1000 = 0.3490$: exactly the loss of part (a), because softmax ignores a common shift.

(c) Mini-batch variance on four examples. Take a model with one parameter and four training examples whose gradients are $g_i = 1, 3, -1, 5$.

1. Full gradient: $(1 + 3 - 1 + 5)/4 = 2$.
2. Per-example variance: $\Sigma = \frac14\big[(1-2)^2 + (3-2)^2 + (-1-2)^2 + (5-2)^2\big] = \frac{1 + 1 + 9 + 9}{4} = 5$.
3. Batches of $B = 2$ drawn with replacement: there are $4 \times 4 = 16$ equally likely ordered pairs. The average of their means is 2 (unbiased), and the variance of their means is $5/2 = 2.5 = \Sigma/B$.
4. Without replacement there are 6 pairs, with variance $\frac{5}{2} \cdot \frac{4 - 2}{4 - 1} = 1.667$.

(d) The learning-rate limit. $L(w) = w^2$ ($\lambda = 2$), $w_0 = 2.4$. With $\eta = 0.3$: $1 - \eta\lambda = 0.4$, so $w = 2.4, 0.96, 0.384, 0.154, \dots$ (converges). With $\eta = 1.1$: $1 - \eta\lambda = -1.2$, so $w = 2.4, -2.88, 3.456, -4.147, \dots$ (diverges).

5. In code #

NumPy and PyTorch check the worked example.

import itertools
import numpy as np
import torch
import torch.nn.functional as F

# (a) softmax cross-entropy and its gradient p - y
z = torch.tensor([2., 1., -1.], dtype=torch.float64, requires_grad=True)
loss = F.cross_entropy(z[None], torch.tensor([0]))      # takes logits
loss.backward()
print(round(loss.item(), 4), z.grad.numpy().round(4))    # 0.349 [-0.2946  0.2595  0.0351]

# (b) shift invariance and log-sum-exp: same loss for z + 998
zb = np.array([1000., 999., 997.])
m = zb.max()
print(round(m + np.log(np.exp(zb - m).sum()) - zb[0], 4))   # 0.349
print(torch.log_softmax(torch.tensor([0., -110.]), 0))      # tensor([   0., -110.])
print(torch.log(torch.softmax(torch.tensor([0., -110.]), 0)))  # tensor([0., -inf]) in float32

# (c) mini-batch variance: Sigma / B with replacement
g = np.array([1., 3., -1., 5.])
means = [(g[i] + g[j]) / 2 for i, j in itertools.product(range(4), repeat=2)]
print(g.mean(), g.var(), np.mean(means), np.var(means))     # 2.0 5.0 2.0 2.5

# (d) gradient descent on w^2 (lambda = 2)
for eta in (0.3, 1.1):
    w, ws = 2.4, []
    for _ in range(4):
        ws.append(round(w, 3)); w -= eta * 2 * w
    print(eta, ws)                     # 0.3 [2.4, 0.96, 0.384, 0.154] / 1.1 [2.4, -2.88, 3.456, -4.147]

# the softmax-before-CrossEntropyLoss floor: log(1 + (K - 1)/e)
print(round(F.cross_entropy(torch.eye(10), torch.arange(10)).item(), 3))   # 1.461

The last line passes perfect one-hot "probabilities" as logits: even then the loss is $1.461$ for $K = 10$, not 0.

6. Where it shows up #

• Language models. GPT-style models are trained with softmax cross-entropy over the vocabulary to predict the next token (lesson 08). The initial loss is about $\log V$: $\log 50{,}257 \approx 10.8$ for the GPT-2 vocabulary. Large-scale runs watch exactly the loss curves of section 3.7 (lesson 12).
• Image classification. ImageNet models use cross-entropy over 1000 classes (initial loss $\log 1000 \approx 6.9$), mini-batches of 256 or more, and SGD with momentum or AdamW.
• Regression and generative models. Diffusion models are trained with a mean squared error on the noise; it is a Gaussian likelihood in disguise.
• ML for science. Physics-informed networks minimize a sum of squared residuals of a PDE at sampled points; a mini-batch of collocation points is a mini-batch in the sense of this lesson. Neural surrogate models of simulations are trained with MSE, or with a Gaussian negative log-likelihood when the network also predicts the noise level $\sigma(x)$.
• Choosing hyperparameters. Learning rate, batch size and the number of steps are chosen on the validation set; the stability limit $2/\lambda$ explains why the best learning rate is often just below the value where training diverges.

7. Common mistakes #

• Softmax (or sigmoid) before the loss. CrossEntropyLoss and BCEWithLogitsLoss take logits. Passing probabilities applies the nonlinearity twice and puts a floor on the loss.
• Forgetting opt.zero_grad(). Gradients accumulate across steps, and the updates grow until training becomes unstable.
• "model.eval() turns off gradients." It does not; it changes dropout and BatchNorm. Use torch.no_grad() for that.
• Choosing hyperparameters on the test set. Then the test loss is no longer an honest estimate. Use the validation set, and the test set once.
• Reading the mini-batch loss as the training loss. The mini-batch loss is noisy; average it over many steps, or compute the loss on the whole training set, before you draw conclusions.
• Sigmoid + MSE for classification. It works, but learns very slowly from confident mistakes. Use cross-entropy.
• "A larger batch is always better." Its noise falls only like $1/\sqrt{B}$ while its cost grows like $B$, and the learning rate usually has to change with it (lesson 12).

8. Check yourself #

1. A 1000-class model starts with logits near 0. What loss do you expect at step 0?
   Answer

   $\log 1000 \approx 6.91$. Near-zero logits give $p_c \approx 1/1000$ for every class.
2. Show that the softmax does not change when you add the same number to every logit.
   Answer

   $e^{z_c + a}/\sum_j e^{z_j + a} = e^a e^{z_c}/(e^a \sum_j e^{z_j}) = e^{z_c}/\sum_j e^{z_j}$. This is why we may subtract $\max_j z_j$ before exponentiating.
3. A binary classifier outputs the logit $z = 3$ for an example with $y = 0$. Compute the BCE gradient $\partial \ell / \partial z$, and the gradient of $\frac12(s(z) - y)^2$.
   Answer

   $s(3) = 0.9526$. BCE: $s - y = 0.9526$. MSE after a sigmoid: $(s - y)s(1 - s) = 0.9526 \cdot 0.9526 \cdot 0.0474 = 0.0430$, about 22 times smaller.
4. Per-example gradients have total variance $\operatorname{tr}\Sigma = 8$. What is the expected squared distance between a mini-batch gradient with $B = 32$ (independent draws) and the full gradient?
   Answer

   $\mathbb{E}\|g_{\mathcal{B}} - \nabla \hat R\|^2 = \operatorname{tr}\operatorname{Cov}(g_{\mathcal{B}}) = \operatorname{tr}\Sigma / B = 8/32 = 0.25$, so the typical distance is about $0.5$.
5. A quadratic loss has Hessian eigenvalues $0.5$ and $10$. Which learning rates make gradient descent converge?
   Answer

   $0 < \eta < 2/\lambda_{\max} = 2/10 = 0.2$. The flat direction ($\lambda = 0.5$) would allow $\eta < 4$, but the steep one sets the limit.
6. Your training loss falls to 0.01 while the validation loss rises from 0.3 to 0.8. What is happening, and name two remedies.
   Answer

   Overfitting: the model fits noise in the training set. Remedies: stop early (keep the parameters with the lowest validation loss), get more data, use regularization such as weight decay or dropout, or use a smaller model.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. What is the gradient of softmax cross-entropy with respect to the logits $z$?

   Show answer

2. Which loss does maximum likelihood give for regression with Gaussian noise of fixed variance?

   Show answer

3. A $K$-class classifier outputs logits near 0. What is its cross-entropy loss?

   Show answer

4. Why does a sigmoid followed by squared error learn slowly from confident mistakes?

   Show answer

5. Why do PyTorch's CrossEntropyLoss and BCEWithLogitsLoss take logits instead of probabilities?

   Show answer

6. Write the log-sum-exp trick.

   Show answer

7. How does the variance of a mini-batch gradient depend on the batch size $B$ (independent draws)?

   Show answer

8. For which learning rates does gradient descent on $\frac{\lambda}{2}w^2$ converge?

   Show answer

9. Name the five lines inside the PyTorch training loop, in order.

   Show answer

10. What does opt.zero_grad() prevent?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Why cross-entropy and not MSE for classification?" — It is the maximum-likelihood loss for a categorical output; its gradient with respect to the logits is $p - y$, which stays large when the model is confidently wrong. With a sigmoid, MSE has gradient $(s - y)s(1 - s)$, which vanishes there.
• "Derive the gradient of softmax cross-entropy." — $\ell = -z_{c^*} + \log\sum_j e^{z_j}$; the derivative of log-sum-exp is $p$, so $\nabla_z \ell = p - y$.
• "What is the expected initial loss of your classifier?" — $\log K$ with small logits. Check it before a long run.
• "Why mini-batches? What happens when you change the batch size?" — Cheaper steps with an unbiased gradient; the variance falls like $1/B$, the cost grows like $B$. The learning rate usually has to be tuned again (lesson 12).
• "Your loss is NaN after 200 steps. What do you check?" — The learning rate (stability limit), the loss computed from probabilities instead of logits, divisions or logs of zero in the data pipeline, and the gradient norm before the explosion.
• "Difference between model.eval() and torch.no_grad()?" — Layer behaviour (dropout, BatchNorm) against graph recording. Evaluation code uses both.

10. Practice #

The notebook for this lesson runs in Colab. It uses NumPy for the first three exercises and PyTorch (on the CPU) for the last two:

1. log_softmax and cross_entropy in NumPy with the log-sum-exp trick, checked against PyTorch on logits as large as 1000.
2. The gradient $p - y$, checked by central finite differences.
3. Measure the variance of mini-batch gradients for several batch sizes and fit the slope $-1$ on a log-log plot.
4. Write the PyTorch training loop for a two-layer MLP on make_moons and reach more than 95% validation accuracy.
5. Find the bug: a loop without zero_grad(). Predict the effect, run it, and fix it.

11. Watch and read #

Gradient descent, how neural networks learn | Deep Learning Chapter 23Blue1Brown · 21 min

The cost function as an average over training examples, and gradient descent as small downhill steps.

The spelled-out intro to neural networks and backpropagation: building microgradAndrej Karpathy · 146 min (watch from 1:51:04 to 2:14)

A dataset, a loss, and the training loop written by hand, including the forgotten zero_grad bug.

Lecture 3 | Loss Functions and OptimizationStanford CS231n (2017) · Justin Johnson · 75 min

Softmax cross-entropy, its expected initial value, and mini-batch gradient descent.

• Simon Prince, Understanding Deep Learning (free PDF): chapter 5 (loss functions from maximum likelihood) and chapter 6 (gradient descent, SGD, Adam).
• Goodfellow, Bengio and Courville, Deep Learning, chapter 6.2 (cost functions and output units) and chapter 8 (optimization), free online.
• Andrej Karpathy, A recipe for training neural networks (2019): the sanity checks of section 3.7 in practice.
• Stanford CS231n notes, Neural networks part 3: learning and evaluation: gradient checks, loss curves, and learning rates.
• PyTorch documentation, CrossEntropyLoss: what the input is expected to be.