Training neural networks

2025 Q2 asked for “an objective function for a neural network”. This stage gives you that, backprop by hand, the classification version with softmax, and the PyTorch training loop from the lab (including the lab’s double-softmax bug).

1. The objective, and why it’s hard to minimise

For a numeric output, minimise the squared error over all parameters \(\theta\) (all the \(w\)’s and \(\beta\)’s):

\[\mathrm{RSS}(\theta) = \sum_{i=1}^{n}\big(y_i - f(x_i)\big)^2 \quad\text{or the MSE}\quad \frac1N\sum_{n=1}^{N}\big(Y_n - f(X_n)\big)^2.\]

The \(w\)’s sit inside \(g\), so there is no closed-form solution: use gradient descent. And unlike linear regression’s single bowl, the objective is non-convex: many valleys, so gradient descent can stop in a local minimum, and where it ends depends on where it starts. One reason: swapping two hidden units gives exactly the same \(f\), so every solution has copies.

Starting point: small random weights, not zeros. With all-zero weights every hidden unit computes the same thing and gets the same gradient, so they stay identical forever. Small, because large weights push sigmoids into their flat tails, where the slope is nearly 0.

2. Backpropagation by hand

The gradient of a sum is the sum of gradients, so work with one point, \(R_i = (y_i - f(x_i))^2\). Write \(z_{ik}\) for unit \(k\)’s weighted sum and \(A_{ik} = g(z_{ik})\). Chain rule:

\[\frac{\partial R_i}{\partial\beta_k} = -2(y_i - f(x_i))\,A_{ik},\qquad \frac{\partial R_i}{\partial w_{kj}} = -2(y_i - f(x_i))\cdot\beta_k\cdot g'(z_{ik})\cdot x_{ij}\]

(with \(A_{i0} = 1\) for \(\beta_0\), and \(x_{i0} = 1\) for the bias \(w_{k0}\)). Both start with the residual: a forward pass computes the \(z\)’s, \(A\)’s and \(f\); then the error is sent backwards, to each output weight scaled by its unit’s activation, and further back to each hidden weight scaled by \(\beta_k\), the slope \(g'\) and the input. Hence backpropagation. With ReLU, \(g'(z) = 1\) if \(z > 0\) and 0 otherwise: a unit that’s off gets no gradient (a unit off for every point is dead).

3. Classification networks: softmax and cross-entropy

Two classes: \(P(Y = 1 \mid X) = \sigma(f(X))\), trained with week 3’s negative log-likelihood. Logistic regression is this with no hidden layer.

\(M\) classes: \(M\) outputs \(f_0, \dots, f_{M - 1}\) sharing the hidden layer, turned into probabilities by softmax, \(p_m = e^{f_m}/\sum_{m'}e^{f_{m'}}\). Encode the label as one-hot (\(y_{im} = 1\) for the true class). The loss is the cross-entropy:

\[\ell(\theta) = -\sum_{i=1}^{n}\sum_{m}y_{im}\log p_m(x_i) = -\sum_i\log p_{y_i}(x_i):\]

minus the log-probability of the true class. Its gradient at the outputs is simply \(p_m - y_{im}\) (predicted minus truth), then backprop as before. Predict the argmax. Parameters: \(K(p + 1) + M(K + 1)\).

4. Fighting overfitting, and the PyTorch loop

A network has many parameters, so it can memorise a small training set. Two strategies (slide 43):

1. Regularisation: add a ridge penalty, \(R(\theta) = \mathrm{RSS} + \lambda\|\theta\|^2\) (not on the biases), \(\lambda\) chosen on validation data. PyTorch’s weight_decay.
2. Slow learning + early stopping: small random start, small steps, watch the validation loss each epoch and stop (keep the best epoch’s weights) when it starts rising while the training loss keeps falling.

Also: fewer hidden units. Minibatches: each step uses a random batch (e.g. 32 points) instead of all \(n\): stochastic gradient descent; one epoch is one pass through the data (\(\lceil n/\text{batch}\rceil\) steps). The lab’s Adam is gradient descent with automatic per-parameter step sizes.

for X_b, y_b in train_loader:      # one epoch, minibatch by minibatch
    out = model(X_b)               # forward pass (logits!)
    loss = criterion(out, y_b)     # CrossEntropyLoss for classes, MSELoss for numbers
    optimizer.zero_grad()          # clear old gradients (they accumulate otherwise)
    loss.backward()                # backpropagation
    optimizer.step()               # theta <- theta - eta * grad (or Adam's version)

Stage test