Backpropagation & Optimization in a Nutshell

"I never see the whole landscape. I only feel the slope under my feet and take one careful step downhill. A few million steps later, people call it intelligence."

A Stochastic Gradient, Stepping Blindly Downhill

In Section 18.1 we set the XOR network's weights by hand. That does not scale past four data points and two hidden units. We need a procedure that, given a measure of how wrong the network is, automatically computes which direction to move each of its (possibly billions of) parameters to make it less wrong. That procedure is backpropagation feeding gradient descent. Both are simpler than their reputation, and understanding them once means that for the rest of Part III you can call loss.backward() and optimizer.step() knowing exactly what they do.

1. The Loss Function: A Scalar to Descend Beginner

Optimization needs a single number to minimize. For classification that number is the cross-entropy loss, which measures the distance between the predicted probability distribution $\mathbf{p}$ and the true class. If the true class is $y$, the loss for one example is the negative log-probability the model assigned to the correct answer:

$$ \mathcal{L} = -\log p_y = -\log \frac{e^{z_y}}{\sum_k e^{z_k}}. $$

A confident correct prediction ($p_y \to 1$) gives loss near zero; a confident wrong prediction ($p_y \to 0$) gives a large loss. The full objective averages this over a batch. The gradient of cross-entropy with respect to the logits has a famously clean form: $\partial \mathcal{L} / \partial z_c = p_c - \mathbb{1}[c = y]$, the predicted probability minus the one-hot truth. Read intuitively, this is just "prediction minus target". Each logit is pushed down in proportion to how much probability the model wrongly placed on that class, and the true class is additionally pulled up. A perfect prediction yields a zero gradient and no further change. This simplicity is why softmax and cross-entropy are always paired. The regression analogue is the mean squared error (MSE), the average of the squared differences between prediction and target, $\frac{1}{n}\sum_i (\hat{y}_i - y_i)^2$; it plays the same role for continuous targets and is the same fidelity measure that underlies the PSNR image-quality metric of Chapter 1.

2. The Computation Graph Beginner

A network's forward pass is a sequence of elementary operations: multiply, add, apply ReLU, take softmax, compute the log. We can draw this as a directed graph whose nodes are operations and whose edges carry values forward. Backpropagation walks this graph in reverse, and to do so it needs each node to know its local derivative, how its output changes when its input changes. Figure 18.2.1 shows the graph for a two-layer MLP and the two passes that flow through it.

3. Backpropagation Is the Chain Rule Intermediate

The chain rule says that if $\mathcal{L}$ depends on $a$ which depends on $w$, then $\partial \mathcal{L}/\partial w = (\partial \mathcal{L}/\partial a)(\partial a/\partial w)$. Backpropagation applies this rule node by node, in reverse topological order, reusing the gradient computed for downstream nodes (the upstream gradient) so no derivative is recomputed. For a linear node $\mathbf{a} = W\mathbf{x}$ with upstream gradient $\mathbf{g} = \partial \mathcal{L}/\partial \mathbf{a}$, the local rules are

$$ \frac{\partial \mathcal{L}}{\partial W} = \mathbf{g}\,\mathbf{x}^\top, \qquad \frac{\partial \mathcal{L}}{\partial \mathbf{x}} = W^\top \mathbf{g}, $$

and for a ReLU node the local derivative is the mask "1 where the input was positive, 0 elsewhere". The code below implements the full forward and backward pass of the one-hidden-layer MLP from Section 18.1 in NumPy, with no autograd, to make every gradient explicit.

# Full forward AND backward pass of a one-hidden-layer MLP in pure NumPy, no
# autograd: the forward computes the cross-entropy loss, and the six backward
# lines apply the chain rule to produce gradients for every weight and bias.
import numpy as np
rng = np.random.default_rng(3)
batch, d, H, C = 16, 20, 32, 4
X = rng.standard_normal((batch, d))
y = rng.integers(0, C, size=batch)

W1 = rng.standard_normal((H, d)) * np.sqrt(2.0 / d); b1 = np.zeros(H)
W2 = rng.standard_normal((C, H)) * np.sqrt(2.0 / H); b2 = np.zeros(C)

# ---- forward ----
z1 = X @ W1.T + b1            # pre-activation, (batch, H)
a1 = np.maximum(0.0, z1)      # ReLU activation
z2 = a1 @ W2.T + b2           # logits, (batch, C)
z2 -= z2.max(1, keepdims=True)
p  = np.exp(z2); p /= p.sum(1, keepdims=True)
loss = -np.log(p[np.arange(batch), y] + 1e-12).mean()

# ---- backward ----
dz2 = p.copy(); dz2[np.arange(batch), y] -= 1; dz2 /= batch   # softmax+CE gradient
dW2 = dz2.T @ a1; db2 = dz2.sum(0)
da1 = dz2 @ W2
dz1 = da1 * (z1 > 0)          # ReLU local derivative
dW1 = dz1.T @ X; db1 = dz1.sum(0)
print(round(loss, 4), dW1.shape, dW2.shape)   # e.g. 1.5231 (32, 20) (4, 32)

Those six backward lines are the entire learning machinery, and a key practical point is hidden in them: the backward pass needs the forward activations ($a_1$, the ReLU mask). This is why training a network costs roughly twice the memory of running it: the intermediate activations must be stored for the backward pass. Chapter 21 returns to this when memory becomes the binding constraint.

4. Gradient Descent and Its Stochastic Cousin Beginner

With the gradient $\nabla_\theta \mathcal{L}$ in hand, the update rule is to step opposite the gradient by a learning rate $\eta$:

$$ \theta \leftarrow \theta - \eta \, \nabla_\theta \mathcal{L}. $$

Computing the gradient over the entire dataset (batch gradient descent) is accurate but expensive and rarely worth it. Stochastic gradient descent (SGD) instead estimates the gradient from a small random mini-batch of examples, trading a noisier gradient for vastly more updates per unit of compute. The noise is not purely a nuisance: it helps the optimizer escape sharp, poorly generalizing minima and saddle points. The learning rate $\eta$ is the single most consequential hyperparameter in all of deep learning: too large and the loss diverges, too small and training crawls or stalls in the first bad valley it finds. Figure 18.2.2 illustrates the three regimes on a one-dimensional loss.

5. Momentum, Adam, and the Optimizer Family Intermediate

Plain SGD struggles in ravines, valleys that are steep in one direction and shallow in another, where it zig-zags. Momentum fixes this by accumulating a velocity, an exponentially decayed running average of past gradients, and stepping along it: $\mathbf{v} \leftarrow \mu\mathbf{v} - \eta\nabla\mathcal{L}$, then $\theta \leftarrow \theta + \mathbf{v}$. The velocity damps oscillation across the ravine and builds speed along its floor. Adam goes further, keeping per-parameter running averages of both the gradient (first moment $\mathbf{m}$) and its square (second moment $\mathbf{u}$), and scaling each parameter's step by the inverse square root of its second moment:

$$ \mathbf{m} \leftarrow \beta_1 \mathbf{m} + (1-\beta_1)\mathbf{g}, \quad \mathbf{u} \leftarrow \beta_2 \mathbf{u} + (1-\beta_2)\mathbf{g}^2, \quad \theta \leftarrow \theta - \eta\,\frac{\hat{\mathbf{m}}}{\sqrt{\hat{\mathbf{u}}}+\epsilon}, $$

where $\hat{\mathbf{m}}, \hat{\mathbf{u}}$ are bias-corrected versions. The per-parameter scaling means rarely-updated parameters get larger effective steps. That forgives a badly tuned base learning rate, which is why Adam is the default first choice in 2026. Its modern variant, AdamW, decouples weight decay (L2 regularization) from the adaptive step, fixing a subtle interaction that hurt generalization; this is the optimizer the training loop in Section 18.5 reaches for. We build up to that scaling in two steps: first the snippet below implements momentum SGD in a few lines to show there is no mystery in optimizer.step(), then the worked case after it traces how a common and a rare feature move under SGD versus Adam, making the forgiveness concrete.

# Momentum SGD from scratch: each parameter gets a velocity buffer that
# accumulates a decayed running average of its gradients, and the update steps
# along that velocity. This is exactly what optimizer.step() does internally.
import numpy as np

class MomentumSGD:
    def __init__(self, params, lr=0.05, mu=0.9):
        self.lr, self.mu = lr, mu
        self.vel = [np.zeros_like(p) for p in params]   # one velocity per param tensor

    def step(self, params, grads):
        for i, (p, g) in enumerate(zip(params, grads)):
            self.vel[i] = self.mu * self.vel[i] - self.lr * g   # accumulate velocity
            p += self.vel[i]                                     # step along it (in place)

# usage sketch with the gradients from subsection 3:
opt = MomentumSGD([W1, b1, W2, b2], lr=0.05, mu=0.9)
opt.step([W1, b1, W2, b2], [dW1, db1, dW2, db2])
print("one optimizer step applied")

6. Learning-Rate Schedules Intermediate

A fixed learning rate is rarely optimal across a whole training run: you want large steps early to make fast progress and small steps late to settle precisely into a minimum. A schedule varies $\eta$ over time. The two dominant modern choices are cosine annealing, which smoothly decays $\eta$ from its peak to near zero following a half-cosine curve, and linear warmup then decay, which ramps $\eta$ up from zero over the first few hundred steps (so the unstable early updates do not blow up) before decaying. Warmup is now standard for training transformers (Chapter 22) and most large models. The schedule interacts with the optimizer: AdamW with cosine decay and a short warmup is the recipe behind a large share of published vision results, and Chapter 21 treats schedule choice as a first-class part of the training recipe.