A network that describes a picture has to carry memory across time. One that sees what a picture means has to run its gradient backward into the pixels. One that invents pictures has to play a two-player game against itself. A3 is four ideas — recurrence, attention, seeing-through-gradients, and generation — and under every one of them is a handful of small, exact functions you can write in numpy.
A companion & practice forge for Stanford's CS 231n Assignment 3 (Image Captioning with RNNs/LSTMs/Transformers, Network Visualization, and Generative Models). 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 (rnn_step_forward/backward, lstm_step_forward, temporal_softmax_loss, scaled dot-product multi-head attention, saliency maps, fooling images, the GAN losses, the SimCLR contrastive loss) are the ones from the homework, shrunk to a scale you can run in your browser with numpy alone — no torch, no GPU, no dataset download.
You feed a photo of a dog on a beach into a captioning model and it emits: "a a the on on a." The picture was understood — a convolutional network already turned it into a rich feature vector. The failure is entirely on the language side: the model is choosing each word with no memory of the words it already chose. Word 3 has no idea word 1 was "a." So it repeats, stalls, and produces nonsense.
A caption is a sequence. To generate one word you must condition on all the words so far. A plain feed-forward layer can't do that — it maps one input to one output and forgets everything. What you need is a network that carries a running summary of the past and updates it one step at a time. That running summary is a hidden state, and the network that maintains it is a recurrent neural network (RNN).
Two little generators emit words one step at a time. The memoryless one picks each word from the image alone — so it loops on the easy words. Toggle memory on and a hidden state (the growing teal bar) carries what was said so far, so each next word depends on the last — a real sentence forms.
The second half of A3 turns the camera around. Instead of using a trained network to produce something, you interrogate it: run its gradient back into the input to see which pixels mattered (saliency), or nudge the pixels to fool it (fooling images). And finally you generate images from scratch with a GAN, and learn representations with no labels at all via SimCLR. Four ideas, one assignment.
A vanilla RNN is one small equation applied over and over. At each timestep it takes the current input xt and the previous hidden state ht−1, mixes them, and squashes the result through a tanh to produce the new hidden state:
Read it as: "take a bit of the new input, add a bit of the running memory, squash to [−1, 1], and that's your updated memory." Wx decides how the input enters, Wh decides how the past carries forward, b is a bias, and tanh keeps the state bounded so it can't blow up.
A3's rnn_step_forward takes a batch of N inputs. The shapes are the whole contract:
| Tensor | Shape | Meaning |
|---|---|---|
x | (N, D) | this step's input for N sequences, each a D-vector |
prev_h | (N, H) | the running memory coming in |
Wx, Wh, b | (D,H), (H,H), (H,) | the shared weights (same at every step) |
next_h | (N, H) | the updated memory going out |
Wx, Wh, b — reused at every step. That's what makes it "recurrent" and lets it handle sequences of any length with a fixed number of parameters. It also means the backward pass has to sum gradient contributions from every timestep.Backprop through one step is a chain-rule exercise. The upstream gradient dnext_h first passes back through the tanh. The derivative of tanh(a) is 1 − tanh(a)², and since next_h = tanh(a), that's just 1 − next_h². So the gradient at the pre-activation is:
From da, everything else is a matrix multiply: dx = da WxT, dprev_h = da WhT, dWx = xT da, dWh = prev_hT da, and db = ∑ da. Forget the (1 − next_h²) factor and every downstream gradient is wrong — the single most common A3 bug.
The tanh curve squashes the pre-activation into [−1,1]. Its derivative 1−tanh² is a bump: near 1 in the middle, near 0 in the saturated tails. Drag the input a — watch the derivative shrink toward the edges. That vanishing bump is exactly why gradients die in deep RNNs (and why LSTMs exist).
rnn_step_backward, why must the upstream gradient be multiplied by (1 − next_h²) before flowing to the weights?The RNN handles vectors, but a caption is words — discrete symbols from a vocabulary. Three small pieces bridge the gap, and together they are the captioning loss you train on.
Each word is an integer index into the vocabulary. Word embedding maps that index to a dense vector by indexing a row of a learned matrix W of shape (V, D) — V words, each a D-vector. Forward is just out = W[x]. The backward pass is the surprising bit: because indexing repeats (the word "the" may appear many times), each gradient must be scatter-added back to its row — np.add.at(dW, x, dout), not a plain assignment, or repeated words overwrite each other.
After the RNN produces a hidden state at each of T steps, a shared linear layer (temporal affine) maps every (N,T,H) hidden state to (N,T,V) vocabulary scores. It's one weight matrix applied at all timesteps at once — a reshape, a matmul, a reshape back.
Now the loss. At each step the model predicts the next word; we compare its V scores against the true next word with softmax cross-entropy. But captions have different lengths, so short ones are padded with a special <NULL> token. We must not count those padded steps in the loss — that's what the mask is for:
The mask is 1 at real words and 0 at <NULL>. Multiply the per-step negative-log-likelihood by the mask before summing, and the padded steps contribute nothing. Forget the mask and the loss is dominated by meaningless padding — a silent, training-wrecking bug.
Each bar is one timestep's negative-log-likelihood in a batch of two captions. The grey steps are padded <NULL>. Toggle the mask: masked steps drop to zero and the loss (number below) falls to the honest value — only real words counted.
temporal_softmax_loss multiply each step's loss by a mask before averaging?The vanilla RNN has a fatal flaw you saw in Chapter 1: its gradient must pass through a tanh at every step, and 1 − tanh² is less than 1, so over many steps the gradient shrinks toward zero. Long-range memory vanishes. The LSTM (Long Short-Term Memory) fixes this with a second state — a cell c — that flows down the sequence on an almost-uninterrupted highway, and three gates that decide what to write, keep, and read.
The LSTM computes four quantities from the same inputs in one shot — a single affine layer produces a (N, 4H) block, sliced into four:
| Gate | Nonlin | Job |
|---|---|---|
| i — input | sigmoid | how much of the new candidate to write into the cell |
| f — forget | sigmoid | how much of the old cell to keep |
| o — output | sigmoid | how much of the cell to read out into the hidden state |
| g — candidate | tanh | the new content proposed for the cell |
The three gates use sigmoid (a soft 0-to-1 valve); the candidate uses tanh (bounded content). Then two update lines:
ct = f·ct−1 + i·g. When the forget gate f is near 1, the old cell passes through almost unchanged — addition, not a squash. Gradient flowing backward along the cell is multiplied by f ≈ 1, not by 1−tanh² < 1. That's why the LSTM remembers across dozens of steps where a vanilla RNN forgets. Drop the f·ct−1 term and you've thrown away the highway; drop the output gate o and the hidden state can't selectively read the cell.One LSTM cell over a few steps. The cell (purple bar) is the memory. Drag the forget gate f: near 1 the cell holds its value across steps (the highway); near 0 it dumps memory each step and behaves like a forgetful RNN. Watch how the hidden state only reads out what the output gate lets through.
Both the RNN and LSTM squeeze the entire past into one hidden vector. That's a bottleneck: to decide word 20 you only get a lossy summary of words 1–19. Attention throws out the bottleneck. Instead of a single running summary, each output position looks directly at every input position and pulls a weighted blend of them — weighting by relevance.
Attention runs on three sets of vectors: queries Q (what each position is looking for), keys K (what each position offers), and values V (the content to pull). A query scores every key by dot product; those scores become weights via softmax; the output is the weighted sum of values:
Read it: QKT is "how much does each query match each key," softmax turns those into a distribution over positions, and multiplying by V pulls the relevant content. The output rows are convex blends of the value vectors.
The dot products grow with the dimension dk: sum dk random products and the magnitude scales like √dk. Large scores push softmax into a near one-hot spike where its gradient nearly vanishes — training stalls. Dividing by √dk keeps the scores at unit scale, so softmax stays soft and gradients flow. Skip the scale and attention is too sharp too early.
Multi-head attention splits the embedding into h chunks, runs a separate attention on each, and concatenates. One head can track subject-verb agreement while another tracks position — different relationships in parallel, then merged. It's the same equation, applied h times on d/h-dimensional slices.
A token-by-token attention heatmap: brighter = query (row) attends more to key (column). Drag the scale: with no 1/√d the scores are large and the map collapses to a hard diagonal spike (each token attends only to one); with the scale it stays soft and spread — gradients survive.
√dk before the softmax?Attention alone has no sense of order — scramble the input tokens and the attention output just permutes with them, because the equation is symmetric in position. And attention alone is linear-ish. The Transformer block adds the two missing pieces: position information, and a per-token nonlinearity, wrapped in the residual + normalization scaffolding that makes deep stacks trainable.
Before attention runs, we add a positional encoding to each token embedding — a fixed pattern of sines and cosines at different frequencies, unique to each position:
Each position gets a distinct fingerprint the network can read off to know "this token is 5th." No learned parameters — just a deterministic pattern added to the embeddings.
A decoder block for captioning runs, in order:
x + sublayer(x)) gives gradients a clean path straight back — the same additive-highway trick the LSTM cell uses, now around each sublayer. Layer norm keeps activations at a stable scale so a deep stack doesn't explode or vanish. Attention does the mixing; the residual and norm make it possible to stack the mixing many times.For captioning, the decoder also uses a causal mask: when predicting word t it may only attend to words ≤ t, never the future — otherwise it would cheat by peeking at the answer. That mask is a triangular −∞ pattern added to the scores before softmax.
The sinusoidal encoding as a heatmap: each row is a position, each column a dimension, running from low frequency (left) to high (right). Every row is a distinct pattern — that's how the network tells positions apart. Drag positions to see more rows; no two are alike.
A3's second half turns a trained, frozen network into an object of study. First question: when the network says "this is a dog," which pixels made it think so? The answer is a saliency map, and it falls straight out of the gradient.
Normally we differentiate the loss with respect to the weights. Here we freeze the weights and differentiate the correct-class score with respect to the input pixels. A pixel whose small change would swing the score a lot is important; a pixel that barely moves the score is not. The saliency at each pixel is the magnitude of that gradient:
Two details that A3 makes you get right: it's the score of the correct class y (not the loss, not another class), and you take the absolute value — a saliency map is a magnitude of influence, so a strongly-negative gradient is just as "important" as a strongly-positive one. (For a color image you also take the max across the RGB channels, but the core is the abs-gradient.)
∂score/∂X — the network never learns anything, its weights never move. You are reading the gradient the network already has, at the input. That's the whole recipe: forward to the class score, backward to the pixels, take abs.Left: a toy input image. Right: its saliency map — brighter where |∂score/∂X| is large. Drag which class you differentiate: the salient region shifts, because different classes care about different pixels. This is the network telling you what it looked at.
Saliency reads the gradient at the input. Fooling images act on it. Take an image the network classifies correctly, pick a wrong target class, and nudge the pixels by gradient ascent on the target's score — over and over — until the network is confident it's the target. The astonishing part: the change to the image is so tiny it's invisible to a human. The picture still looks like a dog; the network is now sure it's an ostrich.
Saliency descended nothing — it just read ∂score/∂X. Fooling images climb that gradient. Each step:
Move the pixels in the direction that raises the target-class score. Two sign traps A3 catches: it must be plus (ascent, raising the target score), not minus (descent, which lowers it and does the opposite of fooling); and the target is the wrong class you're forcing, not the true one. A3 also normalizes the gradient each step (divide by its norm) so the step size is stable regardless of gradient magnitude.
Three class scores as bars. Press step to take a gradient-ascent step toward the target (the rightmost). Watch its bar climb while the image (thumbnail) barely changes. A descent step (wrong sign) would push the target down — the classic bug.
So far the network reads or is read. A Generative Adversarial Network (GAN) makes it invent. Two networks play a game: a generator G turns random noise into fake images, and a discriminator D tries to tell real images from G's fakes. G improves by learning to fool D; D improves by getting harder to fool. At equilibrium, G's fakes are indistinguishable from real.
The discriminator is a binary classifier: label real→1, fake→0. Its loss is binary cross-entropy pushing exactly that:
The generator wants the opposite on the fakes — it wants D to call them real:
Same fake images, flipped label. That's the entire adversarial tension in two lines. Get the labels backwards (the classic bug) and the generator actively makes worse images while thinking it's improving.
D is trained to output 0 on fakes; G is trained to make D output 1 on those same fakes. They pull the same number in opposite directions. As G gets better (its fakes' D-logits rise toward "real"), LG falls — the signal that it's winning. That monotone relationship is exactly what the check verifies.Plain BCE-GANs can get vanishing gradients when D is very confident. LS-GAN swaps cross-entropy for a squared error toward the target labels:
Same structure — real toward 1, fake toward 0, generator pushes fake toward 1 — but the squared penalty gives smoother, non-saturating gradients, so training is more stable. A3 has you implement both.
The fakes' average D-logit on the x-axis (left = D sure they're fake, right = sure they're real). Drag it: the generator loss falls as its fakes look more real; the discriminator's job on those fakes gets harder. The crossover is the adversarial balance point.
A GAN generates; SimCLR represents — and it does so with no labels at all. The idea is beautifully simple: take an image, make two random augmentations of it (crop, color-jitter, flip), and teach the network to give those two views similar embeddings while pushing every other image's views apart. Same image = pull together; different images = push apart. The supervision comes entirely from "these two came from the same photo."
The loss is the normalized temperature-scaled cross entropy (NT-Xent). For a batch, embed every augmented view, L2-normalize, and compute all pairwise cosine similarities divided by a temperature τ. For each view i with positive partner j (its sibling augmentation), the loss is a softmax that treats the positive as the "correct class" among all other views:
It's cross-entropy where the "label" is the positive partner and the "classes" are all other views in the batch.
k ≠ i — crucially excluding i itself. If you leave simi,i (which is 1/τ, huge) in the denominator, it dominates, the fraction collapses toward zero, and the loss reads near-zero no matter what — the model learns nothing. Masking out the diagonal is the one line everyone forgets. With it masked, pulling positives together genuinely lowers the loss.τ sharpens or softens the contrast. Small τ makes the softmax focus hard on the single closest negative (aggressive, can be unstable); larger τ spreads the pressure across all negatives. It's the same temperature knob as any softmax, applied to a similarity distribution.
Four embeddings on a ring: two positive pairs (same image, augmented) joined by a line. Drag alignment: at 0 the pairs are scattered (high loss); slide toward 1 and each pair's two views converge while the pairs stay apart — the NT-Xent loss (below) drops. That descent is SimCLR learning.
The payoff. This is the single most important object of the whole assignment: a live attention matrix over a caption. Watch how, as a Transformer generates each word, it reroutes its attention to the image regions and prior words that matter for that word — and how the 1/√d scaling keeps that routing soft enough to learn.
A token-by-token attention heatmap for a short caption. Press Play and each new query row lights up, attending back over the words emitted so far (a causal, lower-triangular pattern — no peeking at the future). Toggle the 1/√d scale: without it the map snaps to a hard diagonal (each token attends only to itself) and learning would stall; with it the attention stays soft and meaningful.
softmax(QKT/√d). The lower-triangular shape is the causal mask from Chapter 5. The softness of the cells is the 1/√d scaling from Chapter 4. And the whole thing is producing a caption — the sequence problem from Chapter 0, now solved by attention instead of recurrence. This is the machine A3 builds.The debugging reality of A3. Every symptom below points straight at one of the exact functions in this lesson. Read the symptom, find the cause.
| Symptom | Likely cause | Fix |
|---|---|---|
| Gradient check fails by a factor of ~1–2; RNN won't learn | Forgot the tanh derivative in rnn_step_backward | Multiply the upstream grad by (1 − next_h**2) before the matmuls |
| Loss is huge and never drops; captions stay garbage | Missing the mask in temporal_softmax_loss — padded <NULL> steps counted | Multiply per-step NLL by the mask before summing |
| Repeated words get zero gradient in the embedding | Used assignment instead of scatter-add in the embedding backward | np.add.at(dW, x, dout) so repeats accumulate |
| LSTM no better than the vanilla RNN on long sequences | Dropped f·prev_c (no cell highway) or the output gate o | c = f*prev_c + i*g; h = o*tanh(c) |
| Attention collapses to a hard diagonal; training stalls | Forgot to divide scores by √dk | Scale QKT by 1/np.sqrt(head_dim) before softmax |
| Decoder "cheats," perfect train loss, garbage at test | Missing the causal mask — it attended to future words | Add a triangular −∞ mask to scores before softmax |
| Fooling image makes the target less likely | Gradient descent (wrong sign) instead of ascent | Add the gradient: X += lr * grad; differentiate the target score |
| GAN generator produces worse images while "improving" | Flipped labels — generator pushing fakes toward 0 | Generator: BCE(D(fake), 1); discriminator: real→1, fake→0 |
| SimCLR loss stuck near zero; representations don't improve | Left self-similarity in the NT-Xent denominator | Mask the diagonal (set sim[i,i] = −∞) before the sum |
1/√d scale is what keeps the softmax soft enough to train.This lesson completes the CS231n arc after the classifiers and convnets of A1–A2. The ideas here are the seeds of everything modern: attention grows into the Transformer and GPT; the generative game leads to GANs and on to diffusion; SimCLR is the root of contrastive learning and CLIP. Related Gleams: Vision Transformer, Vision-Language Models.