One raw-pixel classifier that will not learn — and the fix is to stop writing one giant tangled function and start building modular layers: each one a tiny box with a forward pass, a cached memory, and a backward pass that hands the gradient back. Snap enough boxes together and you have a fully-connected net; add convolution, pooling, and normalization and you have a ConvNet. Every box you build is exactly one function CS231n Assignment 2 asks you to implement.
A companion & practice forge for Stanford's CS 231n Assignment 2 (Fully-Connected & Convolutional Networks, Batch/Layer Normalization, Dropout). It credits the public course materials and teaches you to implement the core math yourself — it is not a copy-paste answer bank. The real starter-code contracts you meet here (affine_forward/backward, relu_forward/backward, softmax_loss, the SGD-momentum/RMSProp/Adam update rules, batchnorm_forward/backward, layernorm, dropout, conv_forward_naive/backward_naive, max_pool, and spatial_batchnorm) are the ones from the assignment, shrunk to a scale you can run in your browser with numpy alone — no PyTorch, no CIFAR download, no cluster.
You are handed a pile of 32×32 color images and told to build a classifier. So you do the obvious thing: flatten each image into a vector of 3072 numbers, multiply by a weight matrix, and read off ten class scores. It is a single line of matrix algebra. And it works — badly. A one-layer linear classifier tops out well short of what the data allows, because it can only draw straight decision boundaries through pixel space, and cats and dogs do not separate along straight lines.
The fix everyone reaches for is "make it deeper": stack more matrix multiplies, put a nonlinearity between them, and let the network bend its boundaries. But here is the trap that sinks most first attempts. If you write the whole deep network as one giant function — forward pass and gradient tangled together in one block — then the moment you want a fifth layer, or a different nonlinearity, or batch normalization, you are rewriting the derivative of the entire thing by hand. One dropped term in that hand-derived gradient and the network silently fails to learn. There is no error, no crash. It just does not work, and you have no idea which of the forty chain-rule terms you got wrong.
Two classes that are not linearly separable (a ring inside a ring). Drag the slider to add hidden units with a nonlinearity between the layers. At 0 it is a single linear layer — a straight line that can never separate the ring. Add nonlinear hidden units and the boundary bends and the accuracy climbs.
Everything in A2 — the fully-connected net, then the ConvNet, plus batch norm, layer norm, and dropout — is built out of these boxes. Build the boxes correctly and the networks assemble themselves. Get one box's backward pass wrong and you learn to debug it with a numerical gradient check: nudge an input by a hair, watch the output move, and compare that measured slope against the gradient your box returned. If they disagree, your backward pass is wrong. That check is your ground truth throughout this lesson.
Before any specific layer, we agree on the shape of a layer — the contract every box obeys. This is the single most important idea in A2, because once you have it, every layer is just a fill-in-the-blank of the same template.
A layer is a pair of functions. The forward pass takes an input x and any parameters, computes an output, and stashes a cache — whatever it will need later to compute its gradient. The backward pass takes the gradient of the loss with respect to this layer's output (called dout, "upstream gradient") plus that cache, and returns the gradient with respect to each of its inputs.
cachedoutdx (and dw, db if it has params)Why does this compose? Because of the chain rule. The gradient of the final loss with respect to some early variable is a product of local derivatives along the path. Each layer knows only its own local derivative. When you multiply the upstream gradient dout by the local derivative inside backward, you have produced the gradient to hand to the layer below — which does the same thing. The full network's gradient is assembled link by link, and no single function ever needs to know the whole chain.
Read it as: "the gradient reaching my input equals the gradient that reached my output, routed through my own local slope." Every backward pass in this lesson is a concrete instance of that one line.
x and w; batch norm needs the normalized values and the variance. Rather than recompute them, forward stashes them in cache and backward reads them out. Forgetting to cache the right thing is one of the most common A2 bugs: the backward pass then has nothing to route the gradient through.Because layers compose, you can bundle a common pair into a "sandwich." The affine–relu sandwich runs affine forward, then relu forward, caching both; its backward runs relu backward, then affine backward. You will build the pieces; the sandwich is just calling them in order. That is the payoff of the API — new layers cost nothing to combine.
A three-layer chain: affine → ReLU → affine → loss. Press step to push the forward signal through (teal, left to right), then watch the gradient flow back (warm, right to left). Each box multiplies the upstream gradient by its own local slope — the chain rule, assembling itself.
backward(dout, cache) function return?Two boxes power the entire fully-connected net. The affine layer is the linear transform — a matrix multiply plus a bias. The ReLU is the nonlinearity that lets the network bend. Build both, forward and backward, and you can stack a net of any depth.
An affine (fully-connected) layer takes a batch of N inputs, flattens each to a vector, multiplies by a weight matrix w, and adds a bias b:
If each input has D features and the layer has M outputs, then x is [N, D], w is [D, M], b is [M], and out is [N, M]. The reshape(N, −1) is what lets you feed a raw image tensor straight in — it flattens everything after the batch dimension. Forget the + b and every output is shifted; the net can still limp along but it has lost a whole set of free parameters.
The affine layer has three inputs — x, w, b — so backward returns three gradients. Each is the upstream gradient dout (shape [N, M]) routed through the matrix multiply:
| Gradient | Formula | Why |
|---|---|---|
dx | dout · wT, reshaped to x | Each input feature affected every output through w; transpose sends gradient back the way it came. |
dw | xflatT · dout | Each weight w[d,m] multiplied x[:,d]; its gradient sums x weighted by dout across the batch. |
db | dout.sum(axis=0) | The bias adds to every row identically, so gradient sums over the batch. Forgetting this sum is a classic shape bug. |
db sum is the trap. A bias is broadcast across all N examples in the forward pass — the same b is added to every row. The reverse of a broadcast is a sum. So db = dout.sum(0), collapsing the [N, M] upstream gradient down to [M]. Return dout unsummed and you get a shape mismatch (or, worse, a silent broadcast that trains garbage). You will write this exact backward as Kernel 1 in the Studio.The ReLU (rectified linear unit) is almost the simplest useful nonlinearity: relu(z) = max(0, z). Forward, it zeroes every negative and passes every positive through unchanged. Backward, it is a gate: gradient flows through wherever the input was positive, and is blocked wherever the input was negative (or zero).
That indicator 1[z > 0] is why ReLU caches the input z (or just the mask): backward needs to know which units were "on." A unit that was off contributes nothing to the loss, so it gets no gradient. This is also the origin of the dead ReLU problem — a unit stuck negative for every input receives zero gradient forever and can never recover.
The forward curve (teal) is max(0, z) — flat then linear. Drag the input z. When z > 0 the gate is open and the backward gradient passes; when z ≤ 0 the gate is shut and gradient is blocked. The number below is the gradient that would reach the input for an upstream gradient of 1.
db = dout.sum(axis=0) rather than just dout?The affine and ReLU boxes produce class scores — raw, unbounded numbers. To train, we need a single scalar loss that says "how wrong was the whole batch," and its gradient with respect to those scores. That is the softmax loss, and it is the last box before the loss and the first place gradient enters the backward chain.
Softmax turns a row of scores into a probability distribution: exponentiate each score, then divide by the row's sum so the row sums to 1. The cross-entropy loss is the negative log-probability the model assigned to the correct class, averaged over the batch:
The 1/N is the average, and it is easy to drop — do so and your loss is N times too big, which quietly scales every gradient and blows up training. In practice you also subtract each row's max before exponentiating (numerical stability — e of a big number overflows), which does not change the probabilities.
The gradient of the cross-entropy loss with respect to the scores is astonishingly clean:
In words: the gradient at each score is just "the predicted probability, minus 1 at the true class, divided by N." Predict 0.7 for the true class and the gradient there is (0.7 − 1)/N — negative, pushing that score up. Predict 0.2 for a wrong class and the gradient is 0.2/N — positive, pushing it down. Each row of this gradient sums to zero, because probabilities sum to one. Drop the − 1 at the true class and the model has no reason to prefer the correct answer.
L−1 times, a final affine to produce scores, then the softmax loss. Forward runs them in order, caching each. Backward runs the softmax gradient first, then walks the sandwiches in reverse. Add an L2 regularization term — 0.5 · reg · ∑ w² — to the loss and reg · w to each weight's gradient, and you have the complete trainable net. Every piece is a box you have already met.Three class scores for one example whose true class is the first. Drag the true-class score. As it rises above the others, its probability → 1 and the loss → 0; as it falls, the loss climbs steeply. The bars are the softmax probabilities; the number is the loss.
(P − 1[c=y])/N. What is the intuition for the − 1 at the true class?You have a loss and, thanks to the boxes, a gradient for every parameter. The optimizer is the rule that turns that gradient into an update. Plain gradient descent — step opposite the gradient — works but crawls. A2 asks you to implement three smarter update rules, each fixing a specific failure of plain SGD.
Plain SGD stalls in long shallow valleys and rattles across steep narrow ones. Momentum fixes this by keeping a running velocity: each step is a decayed memory of past steps plus the current gradient. Think of a ball rolling downhill — it accumulates speed in a consistent direction and averages out the zig-zag.
The μ · v term (typically μ = 0.9) is the memory. Drop it — keep only −lr · dw — and you are back to plain SGD with no momentum at all.
Some parameters have huge gradients, others tiny; one global learning rate serves neither. RMSProp keeps a running average of each parameter's squared gradient and divides the step by its square root — large-gradient directions get damped, small-gradient directions get amplified:
The √c in the denominator is the whole point — it makes the effective learning rate adapt to each parameter's recent gradient magnitude. The ε (a tiny constant like 1e−8) prevents dividing by zero.
Adam combines both ideas: a momentum-like first moment m and an RMSProp-like second moment v, each an exponential average. Its one subtlety is bias correction: because m and v start at zero, they are biased toward zero in the first steps, so Adam divides each by (1 − βt) to un-bias them:
(1 − βt) correction and the very first Adam steps are far too small — m and v are near zero, so without un-biasing them the update stalls out of the gate and training limps for many iterations. With correction, the first step has the intended magnitude (about lr) immediately. You will write all three update rules as Kernel 4.A narrow, curved loss valley (contours). Press Play to watch four optimizers descend from the same start. SGD zig-zags and crawls; Momentum overshoots but gets there; RMSProp adapts per-axis; Adam combines both. Reset to re-race.
(1 − βt)?Stack six affine–relu layers and try to train them and you hit a wall: the distribution of activations at each layer keeps shifting as the layers below update, so every layer is forever chasing a moving target. This is internal covariate shift, and it makes deep nets painfully slow to train. Batch normalization is the fix — and its backward pass is the most intricate in A2.
For each feature (column), batch norm computes the mean and variance across the batch, normalizes the feature to zero mean and unit variance, then applies a learned scale γ and shift β so the network can undo the normalization if it needs to:
Forget the √(σ²+ε) division and you have only re-centered, not re-scaled — the whole stabilizing effect is gone.
At test time you often have one example, not a batch — there is no batch to compute statistics over. So during training, batch norm also keeps an exponential running average of the mean and variance, and at test time it normalizes using those frozen running stats instead of the (nonexistent) batch stats.
Here is why batch-norm backward is hard. Because μ and σ² are computed across the batch, each output depends on every input in the batch — nudge one input and every normalized output shifts. So the gradient cannot be computed sample-by-sample. The clean form:
where dx̂ = dout · γ. The two subtracted sum-terms are exactly the coupling — they account for how changing one input moves the shared mean and variance. Get them wrong (a common error is to treat μ, σ as constants) and the gradient check fails immediately. The parameter gradients are simpler: dγ = ∑(dout·x̂) and dβ = ∑ dout. You will write forward and backward as Kernel 5.
The raw activations (grey histogram) drift and spread as you drag layer depth. The normalized distribution (teal) stays centered at 0 with unit spread at every depth — then γ, β re-shape it. That stability is what lets gradients flow through many layers.
Batch norm has a weakness: it depends on the batch. With a tiny batch (or a batch of one at test time) its statistics are unreliable. Layer normalization sidesteps this, and dropout is a different tool entirely — a regularizer. A2 asks for both.
Layer normalization is batch norm with the axes swapped. Instead of normalizing each feature across the batch, it normalizes each example across its features:
Because each example is normalized on its own, layer norm is completely independent of batch size — it behaves identically with a batch of 1 or 1000, and needs no running statistics. The implementation is beautiful: it is literally batch norm applied to the transpose. (This is why the assignment can reuse the batch-norm backward by transposing.) The trap is the axis: normalize over axis=0 (the batch) by mistake and you have written batch norm, not layer norm.
Dropout is a regularizer. During training, it randomly sets a fraction of activations to zero, forcing the network not to rely on any single unit — the surviving units must each be useful on their own. A2 uses inverted dropout: it also scales up the surviving units by 1/p (where p is the keep-probability) so that the expected value of the output is unchanged:
The /p is what makes it "inverted," and it is the whole trick: because training already scaled the kept units up, test time needs no scaling at all — dropout is the identity at test time. Forget the /p and the activations shrink by a factor of p at test time relative to training, and the whole net is mis-calibrated.
dx = dout · mask), and is the identity at test time. Every one of them is just another box obeying the forward/cache/backward contract from Chapter 1.A layer of activations, all equal to 1. Drag the keep-probability p. Each frame a random mask drops some units (grey) and scales the survivors by 1/p (tall teal bars). The running average of the output stays near the original value — that is why test time needs no rescaling.
1/p during training?Fully-connected layers throw away the one thing images have going for them: spatial structure. Flattening a 32×32 image destroys the fact that neighboring pixels are related. The convolutional layer keeps that structure by sliding a small learned filter over the image — the defining operation of a ConvNet, and the one whose four-dimensional bookkeeping trips everyone up.
A conv layer takes an input of shape [N, C, H, W] (batch, channels, height, width) and a bank of F filters of shape [F, C, HH, WW]. Each filter slides over the image; at every position it computes a dot product between the filter and the patch of input beneath it, plus a bias:
Two hyperparameters shape the output. Stride s is how far the filter jumps each step; pad is how many zeros ring the input's border (so the filter can center on edge pixels). The output size is (H + 2·pad − HH)/stride + 1 in each spatial dimension. Each filter produces one output channel — the receptive field of an output cell is the patch of input it saw.
Conv backward is the same triple of gradients as affine, but with spatial bookkeeping. For every output cell, the upstream gradient dout[n,f,i,j] flows back to three places:
| Gradient | How it accumulates |
|---|---|
dx (into the input patch) | Add w[f] · dout[n,f,i,j] to the patch of dx under that output cell. Overlapping patches accumulate — a pixel touched by several filter positions sums their contributions. |
dw (into the filter) | Add x_patch · dout[n,f,i,j] to filter f. Because the filter was reused everywhere, its gradient sums over all positions and all examples. |
db (into the bias) | dout[n,f].sum() — the bias touched every output of channel f, so its gradient sums the whole output map. |
The forward is a gather (pull patches in, dot with the filter); the backward is a scatter (push gradient out to every patch the filter touched, accumulating on overlaps). The most common bug is forgetting that overlapping patches must add, not overwrite. You will write forward and backward as Kernel 7.
A 6×6 input (left) and a 3×3 filter. Drag the slider to slide the filter to each stride position; the highlighted patch is dotted with the filter to produce one cell of the output map (right), which fills in as you go. This is the forward pass, one position at a time.
dx from different output positions be added rather than overwritten?Two more spatial boxes finish the ConvNet toolkit. Max pooling shrinks feature maps; spatial batch norm extends batch norm to the four-dimensional world of images. Both are short, and both hide one clean idea.
Max pooling slides a small window (usually 2×2, stride 2) over each channel and keeps only the maximum value in each window. It halves the spatial size, discards precise position (giving a little translation invariance), and has no parameters.
Its backward pass is where the elegance lives. Since the output was the max of a window, only that one winning input affected the output — everything else in the window was irrelevant. So gradient flows back to exactly the argmax position and to nowhere else:
Regular batch norm normalizes each feature over the batch. For images, the natural unit to normalize is a channel: every spatial location of a given channel should share one mean and variance, computed over the batch and all spatial positions. Rather than re-derive batch norm for 4-D tensors, A2 uses a reshape trick:
Collapse everything except the channel axis into one big "batch," run the plain batch-norm you already built, then un-reshape. Each channel is normalized over all N·H·W of its values. This is the whole of spatial batch norm — a shape manipulation wrapped around an existing kernel. The trap is normalizing over the wrong axes (e.g. globally over all channels at once): each channel must have its own statistics. You will write it as Kernel 9.
A 4×4 map pooled with a 2×2 window (stride 2) into a 2×2 output. Toggle backward: forward highlights the max of each window; backward shows the gradient landing on only the argmax cell of each window, zero everywhere else.
Every box you built assembles here into a small ConvNet, trained on a toy dataset, learning to classify in front of you. This is the payoff of the whole assignment: the layers are correct, so the network simply works.
A canonical small ConvNet is: conv → ReLU → max-pool → affine → ReLU → affine → softmax. The conv layer learns spatial filters; ReLU adds nonlinearity; pooling downsamples; the affine layers mix everything into class scores; softmax turns them into a loss. Every arrow is one of your boxes, forward and backward. Nothing else is needed.
Press Play to train. The loss (top curve) falls and the accuracy (bottom curve) climbs as the ConvNet learns to classify two visually distinct toy classes. The small grid shows a learned conv filter sharpening from noise into a detector. Toggle a broken layer to watch training stall — every box has to be right.
/N and the loss stalls or diverges — the network cannot learn from a wrong gradient. That is the whole reason A2 makes you gradient-check every layer: a single wrong backward pass silently kills training, and only the numerical check tells you which box is at fault.And the mechanism is exactly the boxes from every prior chapter: the affine and ReLU from Chapter 2, the softmax loss from Chapter 3, the optimizer from Chapter 4, and the convolution and pooling from Chapters 7 and 8. Assemble correct boxes and the chain rule — and the learning — take care of themselves.
The debugging reality of building layers by hand. When a network won't learn, the symptom usually points straight at one box's forward or backward pass. Read the symptom, find the cause.
| Symptom | Likely cause | Fix |
|---|---|---|
| Gradient check fails on a layer (analytic ≠ numerical) | A wrong term in that layer's backward pass — forgot a sum, a transpose, or the batch coupling | Trust the numerical check; re-derive the local gradient of that one box |
Loss is N× too large; gradients explode | Softmax used sum instead of mean, or forgot /N on the gradient | Average the loss over the batch; divide the score gradient by N |
| Shape mismatch or silent broadcast in affine backward | db not summed over the batch — returned [N,M] instead of [M] | db = dout.sum(0) — the reverse of the forward bias broadcast |
| Half the ReLU units never activate; net underperforms | Dead ReLUs — units stuck negative get zero gradient forever (often too-high learning rate) | Lower the learning rate; use better init; the backward mask is correct, the units died |
| Great train accuracy, terrible test accuracy with batch norm | Using batch stats at test time, or never updating the running mean/var during training | Train uses batch stats + updates running avg; test uses running avg only |
| Dropout net is fine in train, off at test (or vice versa) | Forgot the /p inverted scaling, or applied dropout at test time | Scale by 1/p in train; dropout is the identity at test |
| Conv gradient check fails only with stride < filter size | Overlapping patches overwrite instead of accumulate in dx | Use += when scattering gradient into dx, never = |
| Spatial batch norm gives wrong stats | Normalized globally or over the wrong axes instead of per-channel over N·H·W | Transpose to put C last, reshape to [N·H·W, C], run plain BN |
This lesson is the middle of the CS231n assignment set. It builds on Assignment 1 (kNN, SVM, Softmax, a two-layer net) — the loss and gradient you scaled to modular layers here first appeared there — and points toward Assignment 3 (RNNs, attention, network visualization). Related Gleams: The Transformer, Vision Transformer, GPT from scratch, and the training-block lessons on Loss Functions and normalization.