Your first neural network

Slides: Your first neural network

A network is a function with knobs

Forget the brain metaphor for now. A neural network is a function with many tunable constants. A linear layer multiplies the input by a matrix of weights and adds a bias. A nonlinearity such as ReLU keeps stacked layers from collapsing into one big matrix multiply.

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 . They are not probabilities yet; they are the numbers the model gives to the loss function.

Why ReLU matters

If you stack only linear layers, the whole stack is still linear: W2 @ (W1 @ x) can be rewritten as (W2 @ W1) @ x. ReLU is just max(0, x), but that one bend makes depth useful.

import torch

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

That clipping is also why initialization matters. If activations start enormous, values can explode; if they start tiny or mostly negative, useful signal can disappear.

nn.Module is packaging, not magic

PyTorch modules are classes that register parameters and define a forward method. The notebook asks you to rebuild pieces PyTorch normally gives you.

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])

The important contract is the shape: five examples with two features become five examples with three output scores. The uniform_(-bound, bound) init with bound = 1/sqrt(in_features) is exactly what the notebook asks for (it calls this "uniform Kaiming initialization"), applied to both weight and bias.

The training loop in one breath

A training step is forward -> loss -> zero_grad -> backward -> step. This week you can treat backward() as a black box: it computes how each parameter should change. Week 4 opens that box.

Convolutions add one image-specific idea: a small kernel slides over an image looking for the same pattern everywhere. 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.

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 roughly 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.

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? — the one the ARENA notebook says to watch 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 optional now, but it is the world you enter in week 6: what pattern-detectors inside a trained vision network actually look like.