Your first neural network

Slides: Your first neural network

The puzzle: a program that reads handwriting

The task sounds like it needs magic. Given a 28×28 grayscale image of a handwritten digit, output which digit it is. Nobody can write that function by hand. Try to describe, in code, what makes a scrawled "3" a 3 and not an 8. Every rule you write ("two open curves stacked vertically"?) dies on the next person's handwriting.

By the end of this week's session, you will have written a program that does it with over 95% accuracy. The program is about ten lines of ordinary Python. No rules about curves. No if-statements about strokes. It uses two matrix multiplies with a max(0, x) between them, and training finds the numbers in the matrices.

That program is a neural network. Once you see its parts, the magic evaporates and the engineering remains.

A network is a function with knobs

Forget the brain metaphor for now. A neural network is a function with tunable constants. You choose the architecture, the sequence of operations, and training chooses the constants. The dataset is MNIST, the classic collection of 28×28 grayscale handwritten digits that every ML course trains on first. This architecture is the simplest one that works:

Read the pipeline with last week's eyes: (28, 28) flattens to (784,), a linear layer maps that to (100,), a nonlinearity bends it, and a second linear layer maps it to (10,). The result has one score per digit; the biggest score wins. Every stage is a shape move you already know how to read.

A linear layer multiplies the input by a matrix of weights and adds a bias. The weights are the learnable constants. For the 784 → 100 layer, they form a matrix of shape (100, 784), one row per output. The bias is a learnable per-output offset, shape (100,), added after the multiply so each output can shift independently of the input. A nonlinearity such as ReLU keeps stacked layers from collapsing into one big matrix multiply (more on that in a moment).

For a batch of 32 flattened images, the input has shape (32, 784). Multiplying by the transposed weight matrix yields (32, 100): 100 hidden values for each image, with the batch axis untouched. One matrix operation applies the same learned weights to all 32 examples.

This miniaturised forward pass exposes every number:

# no-run
# 01: a linear layer is one line of math
import torch

x = torch.tensor([[2.0, -1.0]])       # one example, two input features
W = torch.tensor([[0.5, 1.0], [-1.0, 0.25], [0.2, 0.2]])
b = torch.tensor([0.1, 0.0, -0.3])
logits = x @ W.T + b
print(logits)
# raw class scores, shape (1, 3)

Those raw class scores are Logits , the numbers the model gives to the loss function. Softmax can convert them into a positive distribution that sums to 1. The Cross-entropy loss used this week accepts logits directly and performs that conversion internally.

For the values above, the logits are [0.1, -2.25, -0.1]. Class 0 wins because 0.1 is the largest score. The gap matters more than the absolute level: adding 10 to every logit leaves the softmax probabilities unchanged. During training, cross-entropy pushes the correct class above its rivals and cares about relative separation between scores.

Read the W matrix row by row: each of the three output scores is a weighted vote over the input features. Row 0 says "score class 0 as 0.5*x[0] + 1.0*x[1] + 0.1." Training chooses those voting weights. The architecture stays fixed while their values change.

Why ReLU matters

Stacking linear layers still produces one linear function. W2 @ (W1 @ x) can be rewritten as (W2 @ W1) @ x. A deep stack with no nonlinearities is a single matrix in a trench coat.

Test the algebra with this two-line experiment. It uses numpy and runs in your browser:

# 02: two linear layers collapse; ReLU breaks the collapse
import numpy as np

rng = np.random.default_rng(0)
W1 = rng.standard_normal((100, 784))
W2 = rng.standard_normal((10, 100))
x = rng.standard_normal(784)

print(np.allclose(W2 @ (W1 @ x), (W2 @ W1) @ x))   # same function?

def relu(v):
    return np.maximum(0, v)

print(np.allclose(W2 @ relu(W1 @ x), (W2 @ W1) @ x))

Expected output: True, then False. The first line proves that a 784→100→10 stack of pure linear layers matches a single 784→10 matrix. The second line inserts max(0, x) and creates a composed function that no single matrix reproduces.

ReLU applies max(0, x). That one bend makes depth useful.

# no-run
# 03: ReLU clips negative values to zero
import torch

hidden = torch.tensor([-2.0, 0.5, 3.0])
print(torch.relu(hidden))
# tensor([0.0000, 0.5000, 3.0000])

Why does clipping help? ReLU lets a hidden unit act like a conditional feature: "this pattern is present this much, or absent." Compositions of these hinges can approximate curved functions that a straight line cannot. Keep the practical fact: linear layers give you votes; nonlinearities let votes depend on other votes.

Picture one hidden unit drawing a boundary through the input space. On one side its weighted sum is negative, so ReLU silences it. On the other side it passes a positive signal to the next layer. A hundred hidden units draw a hundred boundaries, and the second layer combines the active regions. That piecewise behavior lets a small network separate the bent, messy clusters formed by real handwriting.

That clipping also makes initialization matter. Enormous starting activations can explode; tiny or negative activations can erase useful signal. A ReLU that always receives negative inputs outputs zero forever and learns nothing. The initialization exercise starts the network in the productive middle.

nn.Module packages the computation

So far, a "network" is loose tensors and an expression. PyTorch modules package that work in classes that register parameters and define a forward method. For a working programmer, this is a familiar base-class convention. The notebook has you rebuild the pieces that PyTorch provides:

# no-run
# 04: Linear from scratch (what the notebook asks for)
import torch

class Linear(torch.nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        bound = in_features ** -0.5
        self.weight = torch.nn.Parameter(torch.empty(out_features, in_features).uniform_(-bound, bound))
        self.bias = torch.nn.Parameter(torch.empty(out_features).uniform_(-bound, bound))

    def forward(self, x):
        return x @ self.weight.T + self.bias

layer = Linear(2, 3)
print(layer(torch.randn(5, 2)).shape)
# torch.Size([5, 3])

Two details matter. First, torch.nn.Parameter is a plain tensor wrapped in a flag that says "this is learnable; track gradients for it and let the optimizer update it." That wrapper lets model.parameters() find your weights. Second, the shape is the contract: five examples with two features become five examples with three output scores. The notebook asks you to apply uniform_(-bound, bound), with bound = 1/sqrt(in_features), to both weight and bias. It calls this "uniform Kaiming initialization."

Why that bound? Each output sums in_features products. For unit-scale inputs, dividing the weight scale by sqrt(in_features) keeps the sum near unit scale across different layer widths. It is a variance budget: more terms in the sum require smaller terms.

The bias uses the same bound in this exercise because you are matching PyTorch's Linear initialization. The notebook tests that detail, so a mathematically reasonable alternative can still fail the check. Seeded tests make the contract concrete: construct the layer, inspect weight.shape and bias.shape, then compare its forward pass with torch.nn.functional.linear on the same input.

The training loop in one breath

You now have a function with knobs. Training is the part that turns the knobs, and its skeleton never changes:

forward -> loss -> zero_grad -> backward -> step

In words: run the model on a batch of examples (forward); compute the loss, a single number measuring how wrong the outputs are (for classification: cross-entropy, which punishes the model for assigning low probability to the correct label); clear old gradients (zero_grad); compute how each parameter should change to reduce the loss (backward); and nudge every parameter a small step in that direction (step).

Track the shapes through one MNIST batch. x starts as (64, 1, 28, 28), Flatten turns it into (64, 784), and the network returns logits of shape (64, 10). The labels have shape (64,), one integer class per image. Cross-entropy reduces those 640 scores to one scalar loss. After backward(), each parameter receives a gradient with the same shape as that parameter. The optimizer can then update every weight element by element.

Treat backward() as a black box this week: it computes how each parameter should change. Spend the session on the network itself. Week 3 opens step (the optimizer), and week 4 has you rebuild backward from scratch. By the end of week 4, you will understand every part of this loop.

Convolutions add one image-specific idea: a small kernel slides over an image looking for the same pattern everywhere. Think of a Convolution as a grep pattern run at every image position. The network reuses one 3×3 pattern detector across all locations, so it learns "edge" once. A Residual / skip connection adds the input back to the output, out = x + f(x), which makes deep networks easier to train and foreshadows the transformer's residual stream. Plant that flag now: week 5 applies this same x + f(x) trick relentlessly.

When training misbehaves, overfit one tiny batch before launching a full run. Reuse 32 examples for a few hundred updates and watch whether the loss approaches zero. Success proves that the model, loss, gradients, and optimizer can cooperate. Failure narrows the search to your implementation. Print shapes at every module boundary and inspect one parameter before and after step(). Those checks catch shape mismatches and missing updates.

Pair-session guide

Core work is ARENA 0.2 sections 1 and 2: implement ReLU, Linear, Flatten, and a simple MLP, then train it on MNIST and add a validation loop. Stretch work is section 3 (convolutions) and section 4 (assembling a ResNet-34 and copying in pretrained weights); the two "Bonus" sections (feature extraction, convolutions from scratch) are extra credit beyond that. Rotate roles between "shape caller" and "driver"; before each module test, say the input and output shapes out loud.

What you should see

You should finish with the notebook's cross-entropy loss curve falling over three epochs and a validation-accuracy curve climbing from the 10% random-guess baseline to above 95% on MNIST. If you reach the ResNet stretch, your assembled ResNet-34 with copied-in pretrained weights should correctly classify the sample photos in the "verify your model's predictions" exercise. The milestone is simple but real: you trained a neural network from parts you understand.

That 10% → 95% climb answers the opening puzzle. Nobody wrote rules for what a "3" looks like. You wrote a generic knob-covered function and a loop that turns knobs downhill on a wrongness score. Handwriting recognition fell out. The gap between "what you specified" and "what you got" motivates the rest of this course: from week 6 onward, you open trained networks to find out what the knobs learned.

Where to go next

Week 3 opens up the opt.step() line and makes you write the optimizer yourself. Before then, the notebook's own recommended reading:

• 3Blue1Brown's But what is a convolution?, which the ARENA notebook recommends even if you skip everything else.
• What is torch.nn really? by Jeremy Howard rebuilds this week's abstractions from raw tensors, the same journey in the other direction.
• Zoom In: An Introduction to Circuits is an optional preview of week 6 that shows pattern detectors inside a trained vision network.