Foundations through one digit classifier

Prepare with tensors and named axes

Warm up with one tiny image

Before MNIST, use a 2 by 3 grayscale image. A Tensor is an array of values with one or more axes. This tensor has a height axis with two rows and a width axis with three columns.

A shape records the size of each axis. Here, (2, 3) means two rows and three columns.

import torch

image = torch.tensor([[0, 64, 128], [192, 224, 255]], dtype=torch.float32)
print(image.shape, image.dtype)
print(image[0])       # first row; indexing removes the height axis
print(image[:, 1:])  # every row, last two columns
# torch.Size([2, 3]) torch.float32
# tensor([  0.,  64., 128.])
# tensor([[ 64., 128.], [224., 255.]])

dtype names the kind of value stored in the tensor. float32 stores floating point numbers. The colon in a slice means “take every position” on that axis.

PyTorch applies one operation to every selected value. This is a vectorized operation: tensor code replaces a Python loop.

import torch

image = torch.tensor([[0, 64, 128], [192, 224, 255]], dtype=torch.float32)
brighter = (image + 32).clamp(max=255)
print(brighter)
# tensor([[ 32.,  96., 160.],
#         [224., 255., 255.]])

Addition acts on all six pixels. clamp limits values to the largest valid brightness, 255. The shape stays (2, 3).

Read Einstein notation by axis name

Einstein notation gives tensor axes short names. For our image, use h for height and w for width. Give three column weights the name w too:

The shared w axis lines up columns. Leaving w out of the output sums those columns.

The reading rule is: a named axis shared by the inputs lines up matching values for multiplication. If that axis is missing after ->, add the products along it. Axes kept after -> remain in the result.

import torch

image = torch.tensor([[0, 64, 128], [192, 224, 255]], dtype=torch.float32)
weights = torch.tensor([0.25, 0.5, 0.25])
row_scores = torch.einsum("hw,w->h", image, weights)
print(row_scores)
# tensor([ 64.0000, 223.7500])

Read "hw,w->h" as: keep one result for each height position, multiply matching width positions, and sum over width. ARENA also uses einops.einsum, where the same pattern can use full names: einsum(image, weights, "height width, width -> height").

Send one ray through each pixel

Optional preview

This short check tests whether you can transfer shape, broadcasting, and batching ideas. MNIST does not use rays.

A ray starts at an origin O and follows a direction D. A point at distance parameter t is:

point = O + t * D

One camera ray can pass through one image pixel. A grid of pixels needs a grid of directions. With two rows and three columns, directions.shape == (2, 3, 3): height, width, and three coordinates named xyz.

Scroll the diagram sideways on a small screen.

The image axes select a ray. The final axis stores its x, y, and z direction.

import torch

origin = torch.tensor([0.0, 0.0, 0.0])
directions = torch.tensor([
    [[1, -1, -1], [1, -1, 0], [1, -1, 1]],
    [[1,  1, -1], [1,  1, 0], [1,  1, 1]],
], dtype=torch.float32)
points = origin + 2.0 * directions
print(points.shape, points[0, 1])
# torch.Size([2, 3, 3]) tensor([ 2., -2.,  0.])

PyTorch broadcasts the three origin coordinates across all six rays. The same shape-first method scales to many rays and many objects. ARENA 0.1 uses it to practise batched tensor operations while building a renderer. It later tests whether rays hit objects, but you do not need that algorithm here.

For more practice, use the official 0.0 Prerequisites Colab exercise notebook. The full 0.1 Ray Tracing notebook is optional.

Follow one digit through the model

The Week 1 model

In Week 1, you will train a small neural network to recognize handwritten digits from the MNIST dataset.

MNIST images are grayscale and 28 pixels wide by 28 pixels high. The model receives one image with this shape:

image.shape == (1, 1, 28, 28)

The model runs this computation:

(1, 1, 28, 28) image
        ↓ Flatten
(1, 784) pixel values
        ↓ Linear
(1, 100) hidden values
        ↓ ReLU
(1, 100) hidden values
        ↓ Linear
(1, 10) logits

Scroll the diagram sideways on a small screen.

Keep the batch axis and follow the other axes from pixels to class scores.

Each of the ten logits is a score for one digit from 0 to 9. Training changes the model's weights so that the correct digit tends to have the largest score.

The core work uses Sections 1 and 2 of 0.2 CNNs & ResNets. Convolutions and ResNets are stretch work.

Move from the tiny image to an MNIST digit

The tiny image used height and width. An MNIST input adds batch and channel axes, so the model can process many images and distinguish grayscale from other channel layouts. Our image has four axes:

The two ones mean different things. Axis names make the shape readable.

The Shape tells you the size and meaning of each axis. A number alone is not enough. The first 1 means one image, while the second 1 means one grayscale channel.

Indexing with one integer removes the selected axis. image[0] selects the only image and has shape (1, 28, 28). image[0, 0] selects its grayscale channel and has shape (28, 28).

shape = (1, 1, 28, 28)
batch, channels, height, width = shape

print("batch:", batch)
print("channels:", channels)
print("height:", height)
print("width:", width)
print("one image channel:", (height, width))

You should see the four axis sizes followed by one image channel: (28, 28).

Rearrange the pixels with Flatten

The first model operation is Flatten. It combines the channel, height, and width axes into one feature axis:

(1, 1, 28, 28) → (1, 784)

There are 1 × 28 × 28 = 784 pixel values. Flatten changes how those values are arranged. It does not remove or calculate any values. The batch axis stays separate, so the model can process several images at once.

batch, channels, height, width = 1, 1, 28, 28
features = channels * height * width

print((batch, features))
# (1, 784)

In the live notebook, ARENA supplies a Flatten module. Your SimpleMLP will use it before the first linear layer.

Turn pixels into hidden values

A linear layer has learned Weight values and a learned Bias . The first layer receives 784 pixel values and produces 100 hidden values:

pixels.shape == (1, 784)
weight.shape == (100, 784)
bias.shape == (100,)
hidden.shape == (1, 100)

Each hidden value uses one row of 784 weights. The layer multiplies each pixel by its matching weight, adds the 784 products, and then adds one bias value.

einsum uses the same reading rule as "hw,w->h", now with the full axis names from ARENA's supplied Linear implementation:

hidden = einsum(
    pixels,
    weight,
    "batch in_features, out_features in_features -> batch out_features",
)
hidden = hidden + bias

Read the pattern in three steps:

1. batch remains because each image needs its own result.
2. out_features remains because the layer produces 100 hidden values.
3. in_features appears in both inputs but not the output, so the layer sums over all 784 input values.

The next PyTorch example uses four inputs and three outputs so you can see the complete calculation.

import torch

inputs = torch.tensor([[2, 1, -1, 3]], dtype=torch.float32)
weight = torch.tensor([
    [ 1, 0, 1, 0],
    [ 0, 2, 0, 1],
    [-1, 0, 1, 1],
], dtype=torch.float32)
bias = torch.tensor([1, -1, 2], dtype=torch.float32)

outputs = torch.einsum("bi,oi->bo", inputs, weight) + bias
print(outputs)
# tensor([[2., 4., 2.]])

The supplied Linear class wraps its weights and bias in nn.Parameter. This tells PyTorch that they are learned values. It also lets model.parameters() find them when the optimizer updates the model.

Let ReLU bend the computation

The first linear layer is followed by ReLU . ReLU replaces each negative value with zero and leaves each positive value unchanged:

relu(x) = maximum(x, 0)

hidden = [-3.0, -0.5, 0.0, 1.5, 4.0]
after_relu = [max(value, 0.0) for value in hidden]

print(after_relu)
# [0.0, 0.0, 0.0, 1.5, 4.0]

Without ReLU, two linear layers in a row are still one linear operation. ReLU changes the computation between the layers, so the full model can represent functions that one linear layer cannot.

PyTorch models inherit from nn.Module. The __init__ method stores the submodules. The forward method says how data moves through them:

class SimpleMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = Flatten()
        self.linear1 = Linear(28 * 28, 100)
        self.relu = ReLU()
        self.linear2 = Linear(100, 10)

    def forward(self, x):
        x = self.flatten(x)
        x = self.linear1(x)
        x = self.relu(x)
        return self.linear2(x)

This is the structure you will implement in the first core exercise.

Read ten scores as one prediction

The final layer returns a tensor with shape (1, 10). These ten values are Logits . A logit is a raw score for one digit class.

index:   0    1    2    3    4    5    6    7    8    9
logit:  0.2 -0.4  1.1  0.6 -0.8  3.2  0.0  0.9 -0.2  0.3

The largest logit is 3.2 at index 5, so the predicted digit is 5.

Softmax can turn logits into probabilities that add to 1. During training, use the raw logits with Cross-entropy . PyTorch applies the required probability calculation inside F.cross_entropy, so do not apply softmax first.

Cross entropy gives a small loss when the model assigns a high probability to the correct digit. It gives a larger loss when the correct digit receives a low probability.

Train and validate

Send the loss backward

The forward pass computes the ten logits and the loss. While it does this, PyTorch records which operations produced the result.

Then loss.backward() applies the chain rule through those operations in reverse. It computes a Gradient for every learned weight and bias.

A gradient answers a local question. If this parameter increased by a very small amount, how would the loss change? The optimizer uses that information to update the parameter.

You do not need to implement backpropagation yet. ARENA 0.4 is stretch work if you want to build a small automatic differentiation system.

Repeat five operations to learn

One training step has five operations:

Scroll the diagram sideways on a small screen.

logits = model(images)                  # 1. run the model
loss = F.cross_entropy(logits, labels) # 2. measure the error
loss.backward()                        # 3. compute gradients
optimizer.step()                       # 4. update parameters
optimizer.zero_grad()                  # 5. clear gradients

The optimizer updates each parameter using its gradient. The Learning rate controls the size of the update.

PyTorch adds new gradients to the gradients already stored on each parameter. optimizer.zero_grad() clears them before the next batch. Without this step, the next update would use gradients from more than one batch.

steps = {
    3: "compute gradients",
    1: "run the model",
    5: "clear gradients",
    2: "measure the error",
    4: "update parameters",
}

for number in sorted(steps):
    print(number, steps[number])

You should see run, measure, compute, update, and clear in that order.

Check progress on held-out images

Training loss tells us what happened on the batches used to update the model. We also need to check images that are not used for updates.

A validation loop runs the model on the test set and follows four steps:

1. Run the model inside torch.inference_mode() because no gradients are needed.
2. Take argmax across the ten logits to get one predicted digit per image.
3. Compare the predictions with the true labels.
4. Divide the number of correct predictions by the number of test images.

The model changes after every training batch, so the notebook records training loss for every batch. The model does not change during validation, so the notebook records one validation accuracy after each epoch.

The second core exercise asks you to add this validation loop to the supplied train function.

Look ahead to image-aware models

Our SimpleMLP flattens the image before it learns. It no longer has separate height and width axes, so the model does not directly use the fact that nearby pixels form local patterns.

A Convolution keeps the image axes and applies the same small set of weights at every position. This lets one learned operation find a local feature anywhere in the image.

A ResNet is a deeper convolutional network. It adds direct connections around groups of layers, which makes deep networks easier to train.

Sections 3 and 4 of the ARENA notebook explain and implement these ideas. They are stretch work for this course week.

Results checklist

The core notebook path starts at Simple Multi Layer Perceptron in Section 1 and stops after add a validation loop in Section 2.

You will complete two exercises:

1. Implement SimpleMLP and pass tests.test_mlp_module and tests.test_mlp_forward.
2. Add a validation loop that records one accuracy value after each epoch.

When the notebook works:

• The training loss falls over three epochs.
• Test accuracy is above 80 percent after the first epoch.
• The model usually reaches about 95 percent or more after three epochs.

Five checks before the notebook

You are ready when you can answer these without looking back:

1. What does each axis in (1, 1, 28, 28) mean?
2. Why does Flatten produce 784 values for each image?
3. Why is in_features missing from the einsum output?
4. Why should F.cross_entropy receive raw logits?
5. What are the five operations in one training step?

Stretch work

Complete stretch work only after the required prework and core exercises.

1. Implement ReLU and Linear in 0.2 CNNs & ResNets without using the supplied checkpoint answers.
2. Complete Section 3 on convolutions.
3. Complete Section 4 on ResNets.
4. Use 0.0 Prerequisites for more tensor and einsum practice.
5. Use 0.3 Optimization for more detail on weight updates.
6. Use 0.4 Backpropagation to build a small automatic differentiation system.
7. Use 0.1 Ray Tracing for more batched tensor practice.
8. Complete 0.5 VAEs & GANs after the other Chapter 0 work.