PyTorch autograd
PyTorch
Every loss.backward() call runs reverse-mode AD over the computation graph PyTorch built during the forward pass.
python
loss = criterion(model(x), y)
loss.backward()
optim.step()Ch. 13 · Deep Learning Foundations
Backpropagation Through Layers
Deep Learning Foundations
Ch. 13
a=σ(Wx+b)
a neural layer
Training a neural network is just gradient descent on a very deeply nested function. The miracle is that the gradient — through possibly hundreds of layers and billions of parameters — can be computed in time *linear* in the size of the network, using the chain rule applied in reverse. That algorith
⏱ 2min read◇ 2vizs.✓ 5exercises
Before you start
Foundations you'll lean on here
A quick refresher on what this lesson assumes. Tap any card for a plain-language explanation and a tiny worked example.
Concepts you'll lean on
Tap any to see its definition and a worked example.
★Big idea
Backprop is the chain rule applied in reverse — every layer contributes a transpose-weight matvec and a non-linearity-gradient elementwise multiply.
Without this:
Without backprop, neural networks past 2-3 layers are untrainable in practice.
Training a neural network is just gradient descent on a very deeply nested function. The miracle is that the gradient — through possibly hundreds of layers and billions of parameters — can be computed in time *linear* in the size of the network, using the chain rule applied in reverse. That algorithm is backpropagation, and understanding it is the difference between using deep learning and *engineering* deep learning.
A feed-forward network is a composition: , where each layer takes the previous activation, applies a linear map , and a non-linearity . The loss at the end is a scalar function of all the parameters , and we want .
The chain rule says . Define the upstream gradient . Then the entire backward pass is two recursions: and the parameter gradients .
Backprop step — every term labeled
\boldsymbol\delta^{(\ell)}[d_l]=(\mathbf{W}^{(\ell+1)})^\top\boldsymbol\delta^{(\ell+1)}\odot\sigma'(\mathbf{z}^{(\ell)})
upstream gradienttransposed weightnext layer gradientlocal nonlinearity gradient
Two ops per layer: a transposed-weight matvec, then an elementwise multiply.
Loss landscape and gradient direction
Each arrow is at that point — the direction backprop pushes the parameters.
↻
The genius of backprop is that we never form a giant Jacobian. We only ever multiply a *vector* by a Jacobian — a vector-Jacobian product (VJP) — which is cheap because each layer's Jacobian has known structure (diagonal for elementwise non-linearities, for linear layers). This is the reverse-mode autodiff that PyTorch, JAX, and TensorFlow all implement.
Why backprop is reverse-mode autodiff
1/5
- 1The network is a composition.
- 2Chain rule: gradient is a product of Jacobians.
- 3Carry a tangent left-to-right — fast for low-dim input.
- 4Carry a cotangent right-to-left — fast for scalar output (every loss).
- 5Linear in graph size, regardless of parameter count.
Two failure modes haunt backprop in deep networks. Vanishing gradients: if every layer's Jacobian shrinks magnitudes by , exponentially decays going backward, and early layers stop learning. Sigmoid and tanh activations cause this. Exploding gradients: the opposite. Cure: ReLU activations, residual connections, batch normalization, careful initialization, gradient clipping. These tricks are why we can now train 100-layer networks.
Worked example — backprop through one layer: Two-layer net, , , . Forward: , . Upstream , . Then (ReLU's derivative), so . The dead ReLU kills the gradient on its row.
1/2
Manual backprop for a 2-layer net
Python (in browser)
Expected: Loss prints; gradient norms finite and non-zero
main.pypython
Loading editor…
Ready — click Run to execute in your browser.
Python runs entirely in your browser via Pyodide (~6 MB on first Run, cached after).
Forward pass → Backward pass
Same composition, traversed in two directions.
Worked example — gradient with respect to $\mathbf{W}^{(1)}$: . Only the second row updates — the inactive neuron's parameters do not move. Dead ReLUs never recover.
1/3
In the wild
Where this shows up in production
Each card is a real system, model, or technique that rests on this lesson.
JAX jit + grad
JAX
jax.grad differentiates a Python function; jax.jit compiles it to XLA for GPU/TPU. Together they give numpy-style code that runs at compiled speed.
python
from jax import grad, jit
grad_fn = jit(grad(loss_fn))Gradient checkpointing
PyTorchMemory
Save activations only at every k-th layer; recompute the rest during backward. Trades 2× compute for O(sqrt(L)) memory — what lets LLaMA-70B training fit.
Mixed-precision training (AMP)
PyTorchAMP
Forward in fp16/bf16, gradients accumulated in fp32. Cuts memory and accelerates training on Tensor Cores by 2-4× with minimal accuracy loss.
Exercises
Concept Challenge
Put your understanding to the test. Score + streak + speed all count.
Backpropagation Through Layers
5 rounds1000 max pts
Active recall
Recap flashcards
1 / 3
Concept
Backprop is the chain rule applied in reverse
Tap the card to reveal
Confirm you've got it
3 quick questions. Get 2 right to mark this lesson complete.