Handout #4 — Losses, Gradients, and Optimization

We can specify a neural network; now we need to learn its parameters. A loss measures prediction error, an optimization rule uses gradients to improve the parameters, and the chain rule connects derivatives through the network. This handout develops the foundations for analytical backpropagation in Handout #5.

1Learning as optimization

Slides 56–75

1.1  Data, model, loss, objective

Let $D=\{(x_i,y_i)\}_{i=1}^{N}$ be training examples, $f_\theta$ the predictive function, and $\ell(f_\theta(x_i),y_i)$ the loss for one example. The empirical risk is

The loss is a scalar penalty; the objective averages it over the training data. The notation $\arg\min$ describes the desired parameters, not an algorithm or a promise that we can find a global minimum. During an update, the data and targets are fixed and the parameters change.

1.2  Cross-entropy connects back to language modeling

If class $t$ is correct, its one-hot vector has $y_t=1$ and all other entries zero. Therefore

Cross-entropy is the negative log of the correct-class probability, not that probability itself. The minus sign turns maximizing probability into minimizing loss. We use natural logs here: loss is in nats. At $p_t=1$ the loss is 0, and as $p_t\to0$ the loss tends to infinity.

import torch
import torch.nn.functional as F

# Rows are examples; columns are classes. Targets are class indices.
logits = torch.tensor([[2., 1., 0.], [0., 1., 2.]], dtype=torch.float64)
target = torch.tensor([0, 2])
log_probs = logits.log_softmax(dim=-1)
per_example = -log_probs[torch.arange(2), target]
loss = F.cross_entropy(logits, target)
assert torch.allclose(loss, per_example.mean())
print(per_example.tolist(), loss.item())  # both about 0.407606
# Try target = torch.tensor([2, 0]): both losses become about 2.407606.

Pass logits to F.cross_entropy; it computes the log-softmax internally. Taking softmax first gives it different inputs and a different loss. The stable calculation uses $\log p_t=s_t-\log\sum_j e^{s_j}$, evaluating the log-sum-exp after subtracting the maximum logit. This avoids overflowing exponentials. Mathematical softmax positivity does not prevent numerical underflow in a naive implementation.

1.3  Why step against the gradient?

The gradient collects the sensitivities to all parameters. For a small change $\Delta\theta$, the first-order approximation is

Choosing $\Delta\theta=-\eta\nabla J(\theta)$ with $\eta>0$ makes the predicted change $-\eta\|\nabla J(\theta)\|^2$. For a differentiable objective with nonzero gradient, a sufficiently small step decreases the loss. The negative gradient is the steepest local direction under a fixed Euclidean step length; a large finite step can still overshoot.

Switch the widget to Two minima to see why local progress need not find the best possible value. Neural-network objectives are generally nonconvex; the convex bowl explains step size, not a global convergence guarantee for a network.

1.4  Minibatches: estimating the direction

Full-batch gradient descent uses all $N$ examples per update. A minibatch uses $B$ examples; strictly, stochastic gradient descent uses one, though the name commonly includes minibatches. For uniformly sampled examples at fixed $\theta$, the batch-average gradient is an unbiased estimate of the full-data average gradient. One batch can point in a poor direction; unbiased means correct in expectation, not correct on every step.

Smaller batches permit frequent updates with less activation memory, but introduce more noise. Bigger batches often use hardware efficiently. An epoch is one pass over the training set; with $N$ examples, batches of size $B$, and a retained final partial batch, it has $\lceil N/B\rceil$ updates. Noisy updates need not decrease the full-data loss at every step.

1.5  Fitting is not evaluation

Use the training set to update weights, a held-out validation set to choose architecture, learning rate, and stopping point, and a held-out test set for final evaluation after those choices. Repeatedly selecting models by their test scores turns the test set into another validation set.

Low training loss with high validation loss suggests overfitting. High loss on both can reflect insufficient capacity, unsuccessful optimization, or a data/implementation problem. Split text at an appropriate document or source boundary before creating overlapping windows so near-duplicate contexts do not leak across splits.

2Algebra and calculus for backpropagation

Slides 76–89 · Shape bridge to the network equations

2.1  Forward shapes before derivative shapes

Our mathematical input $x$ is a column vector, so one layer computes $Wx+b$. In code we stack examples as rows of $X\in\mathbb R^{B\times d}$:

The bias $b\in\mathbb R^h$ is broadcast to every row. These are two conventions for the same operation, not different models. Each row uses the same parameters.

import torch
import torch.nn as nn

torch.manual_seed(0)
B, d, h, m = 5, 4, 6, 3
x = torch.randn(B, d)
first, second = nn.Linear(d, h), nn.Linear(h, m)
hidden = torch.relu(first(x))
logits = second(hidden)
assert torch.allclose(first(x), x @ first.weight.T + first.bias)
assert hidden.shape == (B, h) and logits.shape == (B, m)
assert sum(p.numel() for layer in (first, second)
           for p in layer.parameters()) == 51
print(hidden.shape, logits.shape)  # (5, 6), (5, 3)

For $X$ of shape $(B,T,d)$, nn.Linear(d,h) maps the last dimension and produces $(B,T,h)$. The matrix work costs $BTdh$ multiply-accumulates. Broadcasting checks compatible sizes, not the intended meaning of an axis: write down and assert shapes.

2.2  Partial derivatives and gradients

A partial derivative changes one argument while holding the others fixed. For the example in slides 78–80, $z(x,y)=y^4x^5$:

The gradient is a column vector in the declared input order $(x,y)$. A derivative is itself a function until we evaluate it at a point.

2.3  A Jacobian records every output–input sensitivity

For $f:\mathbb R^n\to\mathbb R^m$, the Jacobian has outputs as rows and inputs as columns:

For $f(x_1,x_2)=(x_1^2+x_2,\ x_1x_2,\ 3x_2)^\top$,

If $m=1$, the Jacobian is the row $\nabla f^\top$, not the gradient column itself. If $m=n=1$, it is the ordinary scalar derivative. This transpose convention matters in the chain rule.

For a vector output $f(W)\in\mathbb R^m$ and a matrix input $W\in\mathbb R^{n\times p}$, the unflattened derivative has shape $(m,n,p)$ and entries $\partial f_i/\partial W_{jk}$. Flattening $W$ in a fixed order gives the equivalent $m\times(np)$ Jacobian. For a scalar loss, we conventionally store the derivatives with the original parameter shape $(n,p)$.

2.4  Composition: multiply along a path, add across paths

For a scalar chain $x\mapsto y(x)\mapsto z(y)$,

If $y=x^2$ and $z=3y+1$, then $dz/dx=3(2x)=6x$. The local derivative of the outer function is evaluated at the inner function's value.

For $g:\mathbb R^d\to\mathbb R^n$ followed by $f:\mathbb R^n\to\mathbb R^m$, the same rule becomes a matrix product:

The order follows the dependencies and the shapes. Matrix multiplication automatically sums over the intermediate coordinates. In particular, if one scalar $x$ affects several intermediates $y_i$, which all affect a scalar $z$,

2.5  Code lab: check the chain rule numerically

Finite differences estimate the effect of perturbing an input. The following uses only Python's standard library. Predict the row and column shapes before running it.

from math import isclose

def g(x):
    return [x * x, 3 * x]          # R -> R^2

def f(y):
    return y[0] * y[1]            # R^2 -> R

x = 2.0
y = g(x)
J_f = [y[1], y[0]]                # 1 x 2: [dz/dy1, dz/dy2]
J_g = [2 * x, 3.0]               # 2 x 1: [dy1/dx, dy2/dx]^T
contributions = [a * b for a, b in zip(J_f, J_g)]
chain = sum(contributions)
eps = 1e-5
numeric = (f(g(x + eps)) - f(g(x - eps))) / (2 * eps)
assert contributions == [24.0, 12.0]
assert isclose(chain, numeric, rel_tol=1e-7)
print(chain, numeric)             # both approximately 36
# Try eps = 1.0, then 1e-14. Why can either hurt the comparison?

One could perturb every parameter to estimate its derivative, but a central-difference gradient for $P$ parameters costs $2P$ loss evaluations. It is useful for checking tiny examples, not for training a large network. The chain rule gives the mathematical alternative; backprop will organize it so intermediate computations are reused.

3Ready for analytical backprop?

Slide 89 recap · Slide 90 transition

We can now specify the forward computation and the scalar objective. For one labeled example:

Training needs $\nabla_{W_1}\ell$, $\nabla_{b_1}\ell$, $\nabla_{W_2}\ell$, and $\nabla_{b_2}\ell$. Backpropagation computes these gradients using the chain rule; gradient descent uses them to update the parameters. They are distinct jobs.

Show all answers

4Review route

If one part felt unfamiliar, revisit the relevant worked example rather than memorizing all the notation at once:

• For dimensions and matrix multiplication, revisit Handout #1's linear algebra review, then redo §2.1 here.
• For log probabilities, revisit Handout #2's cross-entropy and perplexity discussion, then redo §1.2 here.
• For derivative notation, revisit Handout #1's calculus review, then redo the two-path example in §2.4.