Backpropagation — Derived by Hand on a 2-Layer Net

Why this session exists

Backpropagation computes derivatives of the computation you actually wrote. The optimizer then uses those derivatives to change parameters. A model can train with a wrong gradient, and a correct gradient can accompany a poorly chosen objective or update. We therefore derive the equations, check all weights and biases against finite differences and PyTorch, and test a deliberately broken derivative.

You will derive a two-hidden-layer classifier's backward pass, explain the softmax/cross-entropy simplification, extend the same machinery to four hidden layers, and diagnose gradients using shapes, magnitudes and graph structure. The title's “2-layer” shorthand is ambiguous: this lesson's main network has two hidden layers and three affine parameter layers.

Prerequisites

Use S097 forward passes, scalar chain rules, matrix shapes and S096 activation derivatives. Gradients below are derivatives of a scalar mean loss with respect to row-batch arrays.

(a) Intuition · 5 min

An assembly-line defect can be traced backward through each transformation. In backprop, however, “blame” is a local sensitivity, not moral or causal attribution: it tells how an infinitesimal change in an intermediate value changes this loss. At a branch, contributions from every use of a value must be added.

For $u=x^2$ and $L=u+3u$, the upstream derivative into $u$ is $1+3=4$, so $dL/dx=8x$. Overwriting the first contribution with the second would give $6x$, a valid-looking but wrong number. This is the same accumulation principle used for shared weights and residual branches.

The one sentence to memorise

Start from $dL/dL=1$; in reverse graph order multiply each upstream derivative by the operation's local derivative, summing contributions where paths meet. For $Z=AW+b$ and upstream $G=dL/dZ$:

Then multiply elementwise by the activation derivative. Do not update a weight before it has been used for all derivatives of that forward pass. loss.backward() writes gradients; optimizer.step() changes parameters. Backprop is subject to floating-point rounding and nondifferentiability conventions, not exact real arithmetic in a computer.

A quick history

The conceptual relationship is to reverse-mode automatic differentiation: operations have local derivative rules, and a graph traversal composes them. PyTorch's autograd documentation describes graph reconstruction and saved tensors directly. This mechanism is sufficient to understand modern engines without unsupported claims about who read a particular thesis or one date when deep networks suddenly became trainable.

The whole chain rule as a graph

Each bias derivative is the row sum of its corresponding delta. Notice the corrected input derivative: Delta1 @ W1.T, not inputs.T times the delta. The latter contributes to the weight derivative and has a different shape.

The derivation, step by step

For $Z_{bj}=\sum_i A_{bi}W_{ij}+b_j$, expand one coordinate:

This is exactly $(A^TG)_{ij}$. Likewise, $\partial L/\partial A_{bi}=\sum_jG_{bj}W_{ij}$ gives $GW^T$, and each shared bias appears once per row, giving a sum over rows. If $A$ is $(B,H_1)$ and $G$ is $(B,H_2)$, the three shapes are $(H_1,H_2)$, $(H_2,)$ and $(B,H_1)$ respectively.

For two hidden ReLU layers:

For tanh use $1-A^2$; for sigmoid use $A(1-A)$. ReLU's derivative at exactly zero is conventionally chosen as zero here. Shapes are necessary checks but not sufficient: using the wrong equal-sized activation can pass every shape assertion.

Why (softmax + cross-entropy) has a clean gradient

For one example with normalized target $y$ and logits $z$, write cross entropy directly as

The derivative of the first term with respect to $z_k$ is $e^{z_k}/\sum_j e^{z_j}=p_k$; the second contributes $-y_k$. Thus $d\ell/dz=p-y$. Equivalently, the softmax Jacobian is $J_{jk}=p_j(\delta_{jk}-p_k)$, and multiplying by $-y_j/p_j$ and summing gives the same result. Divide by $B$ for the unweighted batch mean. With unnormalized target mass $s=\sum_jy_j$, the derivative is $sp-y$, so normalized targets are an assumption, not decoration.

For stable computation subtract $m=\max_jz_j$ and use $m+\log\sum_j e^{z_j-m}$. With logits $(1000,-1000)$ and the second class correct, the loss is about 2000 and the gradient is approximately $(1,-1)$. Computing a tiny probability, adding $10^{-12}$, and taking its log instead caps the apparent loss near 27.63; retaining the original $p-y$ gradient would no longer differentiate that altered expression.

Other pairings: sigmoid plus BCE gives $\sigma(z)-y$. Sigmoid plus squared error gives $2(\sigma(z)-y)\sigma(z)(1-\sigma(z))$, which can have a tiny gradient for a confidently wrong prediction. Probability MSE/Brier score is still a legitimate scoring objective, not forbidden mathematics. Linear-output MSE has derivative $2(\hat y-y)/B$, appropriate for a squared-error regression objective but not universally optimal for every continuous-target problem. Batch parameter gradients can cancel even when individual logit gradients are nonzero.

Reverse mode, memory and locality

A Jacobian maps input perturbations to output perturbations. Forward mode computes Jacobian–vector products for input directions; reverse mode computes vector–Jacobian products seeded at outputs. For a scalar loss and many parameters, one reverse traversal yields the whole gradient without materializing a giant Jacobian. Computing every coordinate via separate forward directions would require many sweeps. Forward mode remains useful for few inputs/many outputs or directional derivatives; it is not inherently inferior.

The backward needs selected forward values, which may be saved or recomputed. It need not store every intermediate forever. Activation checkpointing trades extra computation for fewer saved tensors; parameter, gradient, optimizer and activation memory all matter, so there is no universal claim that activation memory causes OOM first.

Where gradients vanish or explode

A sequence of weight matrices and activation Jacobians can contract or amplify perturbations. ReLU does not remove small/large singular values of the weight products. Residual blocks have derivative $I+J_F$, which supplies another path but does not guarantee stable norms. Normalization can change conditioning; it does not “calibrate” probabilities.

Inspect layer-wise activation ranges, nonfinite values, gradient norms and parameter/update norms. A nonfinite input can contaminate downstream computations, but masks and graph structure determine how far it spreads. Gradient clipping rescales a finite gradient $g$ to $g\min(1,c/\|g\|)$; it does not repair NaNs, detached graphs or incorrect targets. Reject or diagnose nonfinite data and gradients before updating, retaining a verified checkpoint for recovery.

The gradient check — a diagnostic, not a proof

A centered difference $(L(\theta+\epsilon e_i)-L(\theta-\epsilon e_i))/(2\epsilon)$ approximates one derivative. Large epsilon has truncation error; tiny epsilon suffers cancellation. Use float64, deterministic data, fixed stochastic masks, several epsilons and both absolute and relative tolerances. Relative error is unstable when both gradients are near zero. Avoid crossing ReLU kinks; inspect the saved masks at perturbed points if a failure is localized there. A successful finite sample does not prove all possible inputs correct.

(c) Hands-on · 25 min

This standalone CPU program supports arbitrary hidden depth. It checks every weight and bias of a tiny two-hidden-layer network, compares NumPy against PyTorch 2.9 float64, verifies the resulting SGD step, trains a bounded generated-data fixture, and executes the four-hidden-sigmoid extension. No GPU or dataset download is used.

import numpy as np
import torch
from scipy.special import expit, logsumexp
from sklearn.datasets import make_classification
 
torch.set_num_threads(1)
def init(sizes,seed=0):
    rng=np.random.default_rng(seed)
    return [(rng.normal(0,.3,(a,b)),rng.normal(.2,.05,b))
            for a,b in zip(sizes[:-1],sizes[1:])]
 
def calc(X,y,p,activation='relu'):
    a=X; acts=[X]; zs=[]
    for i,(W,b) in enumerate(p):
        z=a@W+b; zs.append(z)
        a=(np.maximum(z,0) if activation=='relu' else expit(z)) if i<len(p)-1 else z
        acts.append(a)
    logp=a-logsumexp(a,axis=1,keepdims=True)
    loss=-logp[np.arange(len(y)),y].mean()
    dz=np.exp(logp); dz[np.arange(len(y)),y]-=1; dz/=len(y)
    grads=[None]*len(p)
    for i in reversed(range(len(p))):
        grads[i]=(acts[i].T@dz,dz.sum(0))
        if i:
            derivative=(zs[i-1]>0) if activation=='relu' else acts[i]*(1-acts[i])
            dz=(dz@p[i][0].T)*derivative
    return float(loss),grads,zs
 
def check(X,y,p,activation='relu'):
    loss,grads,zs=calc(X,y,p,activation)
    if activation=='relu': assert min(np.min(abs(z)) for z in zs[:-1])>1e-4
    worst=0.
    for layer,(W,b) in enumerate(p):
        for part,arr in enumerate((W,b)):
            for idx in np.ndindex(arr.shape):
                saved=arr[idx]; analytic=grads[layer][part][idx]
                for eps in [1e-4,1e-5,1e-6]:
                    try:
                        arr[idx]=saved+eps; plus=calc(X,y,p,activation)[0]
                        arr[idx]=saved-eps; minus=calc(X,y,p,activation)[0]
                    finally: arr[idx]=saved
                    numeric=(plus-minus)/(2*eps)
                    assert np.isclose(analytic,numeric,rtol=2e-4,atol=2e-8), (layer,part,idx,analytic,numeric)
                    worst=max(worst,abs(analytic-numeric))
    tp=[(torch.tensor(W,requires_grad=True),torch.tensor(b,requires_grad=True)) for W,b in p]
    a=torch.tensor(X)
    for i,(W,b) in enumerate(tp):
        a=a@W+b
        if i<len(tp)-1: a=torch.relu(a) if activation=='relu' else torch.sigmoid(a)
    tl=torch.nn.functional.cross_entropy(a,torch.tensor(y)); tl.backward()
    assert np.isclose(loss,tl.item(),atol=1e-12)
    for pair,gpair in zip(tp,grads):
        for value,g in zip(pair,gpair):
            assert np.allclose(value.grad.numpy(),g,rtol=1e-9,atol=1e-10)
            before=value.detach().numpy().copy()
            with torch.no_grad(): value.add_(value.grad,alpha=-.1)
            assert np.allclose(value.detach().numpy(),before-.1*g,atol=1e-12)
    print(activation,'layers',len(p),'max finite-difference absolute error',worst)
    return grads
 
rng=np.random.default_rng(2); X=rng.normal(size=(5,3)); y=np.array([0,1,0,1,1])
p=init([3,4,4,2]); grads=check(X,y,p)
# Wrong but shape-compatible activation for output dW: negative control.
_,_,zs=calc(X,y,p); a1=np.maximum(zs[0],0)
logp=zs[-1]-logsumexp(zs[-1],axis=1,keepdims=True)
dz=np.exp(logp); dz[np.arange(5),y]-=1; dz/=5
assert not np.allclose(a1.T@dz,grads[-1][0],rtol=1e-4,atol=1e-8)
q=init([3,3,3,3,3,2],4); check(X,y,q,'sigmoid')
for _ in range(50):
    _,gs,_=calc(X,y,q,'sigmoid')
    for (W,b),(dW,db) in zip(q,gs): W-=.1*dW; b-=.1*db
print('four-hidden sigmoid gradient norms', [float(np.linalg.norm(g[0])) for g in calc(X,y,q,'sigmoid')[1]])
Xt,yt=make_classification(n_samples=180,n_features=4,n_redundant=0,n_informative=4,
                         n_classes=2,random_state=3)
t=init([4,8,4,2],8); initial=calc(Xt,yt,t)[0]
for _ in range(150):
    loss,gs,_=calc(Xt,yt,t)
    for (W,b),(dW,db) in zip(t,gs): W-=.1*dW; b-=.1*db
final=calc(Xt,yt,t)[0]; assert np.isfinite(final) and final<initial
print('training-fixture loss',initial,final)
w=torch.tensor(2.,requires_grad=True)
(w*3).backward(); assert w.item()==2 and w.grad.item()==3
(w*3).backward(); assert w.item()==2 and w.grad.item()==6
x=torch.tensor(2.,requires_grad=True); u=x*x; (u+3*u).backward(); assert x.grad.item()==16
print('branch accumulation and backward-without-step passed')

Anatomy of the script

calc differentiates precisely the same logit loss used by finite differences. The parameter is restored in finally, so a failed perturbation cannot corrupt subsequent checks. Biases are included. The PyTorch arrays use the same (in,out) orientation in explicit matrix multiplication; the framework module transpose is handled in S100.

The negative control replaces the output's saved second hidden activation with the first. Equal widths let its shape pass, but the derivative comparison rejects it. The four-hidden-layer extension has five affine layers, not just an added W4 to a three-affine-layer network. Its printed norms are observations on this fixture, not a guaranteed factor-of-four attenuation or proof that all ReLU networks balance gradients.

The generated-data training section verifies decreasing training objective only. It does not label training accuracy as generalization. To select architecture, learning rate or epoch count for predictive evaluation, add the development/validation/test separation from S095.

(d) Production reality · 15 min

A custom loss can have a factor-of-two bug while its learning curve still declines. Compare against an independently computed reference and numerical differences on multiple inputs; “it trains” is weak evidence. Autograd differentiates the implemented graph, so it too can faithfully compute the wrong intended objective.

For memory pressure, first identify what is allocated. Smaller microbatches reduce live activation memory; accumulation matches a full-batch mean only when per-example computations and weighting match. BatchNorm statistics and stochastic layers can make microbatch execution differ. For unequal microbatch sizes, weight loss sums by the total count instead of averaging each mean equally, and flush a final incomplete group.

Checkpointing recomputes selected activations. Its memory/time trade depends on placement, RNG handling and stateful operations; no universal square-root saving or exact extra-forward cost is asserted. Mixed precision uses different dtypes for different operations. FP16 and BF16 have different exponent ranges, and parameters/gradients may remain float32. Under loss scaling, unscale accumulated gradients before clipping and stepping. This is a source-reviewed mechanism, not an executed accelerator result.

Where this shows up in the rest of the plan

S099 consumes gradients. S100 explains leaf tensors, detached graphs and mode flags. S101 changes the forward map and therefore its derivative. Attention, shared embeddings and adapter training all reuse the same branch-summing rule; which parameters receive updates is a separate choice from which derivatives are needed upstream.

(e) Recall + stretch · 10 min

1. Output gradient? $(P-Y)/B$ for normalized categorical targets and unweighted mean CE.
2. Why $A^TG$? Expand one weight coordinate and sum its effect over examples.
3. Required shape? Each parameter gradient matches that parameter exactly; matching alone does not establish correctness.
4. ReLU mask? Multiply upstream derivatives by 1 on positive pre-activations, 0 on negatives, with zero chosen at the kink.
5. Relative error 0.3? Investigate objective/reduction mismatch, dtype, epsilon, stochasticity and kinks before concluding which derivative is wrong. Use absolute error near zero.
6. Can ReLU gradients vanish? Yes, through inactive branches and contracting weight products.
7. Dead unit? Its ReLU data-gradient is blocked while all relevant pre-activations remain negative; inspect whether this persists across data and updates.
8. Clipping? Rescale finite gradients above a chosen norm bound. Choose the threshold empirically; it does not cure invalid numerics.
9. Autograd versus manual? Same local chain-rule composition, automatically tracked for supported differentiable operations; Python control flow selects the executed graph.
10. Four-hidden sigmoid stretch: Use five affine parameter pairs, sigmoid derivatives in every hidden layer, float64 checks and recorded layer norms. The executable extension demonstrates the procedure without assuming a particular layer ordering of gradient magnitudes on all problems.

Explain-out-loud test

Derive both $dW_2=A_1^T\Delta_2$ and $dA_1=\Delta_2W_2^T$, explain the log-sum-exp derivative, and describe one bug that passes shape checks and one reason a finite-difference check can fail without a derivation error.

What comes next

Optimizers decide what to do with these derivatives; backprop itself makes no parameter step.

Sources and execution boundary

PyTorch 2.9 autograd mechanics, cross entropy, AMP and residual-learning primary paper. All displayed Python is CPU-executed with NumPy/SciPy and PyTorch float64 checks. Checkpointing and mixed-precision discussion is source analysis, not GPU execution or throughput evidence.