Foundations through one digit classifier

Warm up with tensors and a simple ray example

What is a tensor?

A Tensor is a multidimensional array of numbers. When you meet a new tensor, read its shape first. The shape tells you what kind of data you are looking at.

Warm up with one tiny image

The warm-up builds toward two bigger examples: rays, the lines a camera casts through an image plane in ray tracing, and MNIST, a dataset of handwritten digits. Before those, start with a tiny 2 by 3 grayscale image. 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.

Suppose you want to brighten the image by adding 32 to every pixel. In plain Python that means two loops: one over the rows and one over the columns, adding 32 to one pixel at a time. Tensor code needs no loop. PyTorch applies one operation to every selected value at once. This is a vectorized operation. Vectorized code is shorter and much faster, because PyTorch does the work in fast compiled code instead of one Python step per pixel.

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

The addition acts on all six pixels at once. clamp caps every value at 255, the largest valid brightness. Without it, the two brightest pixels would land above 255. The shape stays (2, 3).

Cast rays through a pixel grid

In this course, ray tracing is practice with tensors. You cast one ray per pixel and find a point along each ray. ARENA's full exercise builds a renderer, a program that turns a 3D scene into a 2D image. You do not need that this week, only the shape skills.

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

point = O + t · D

One camera ray passes through each 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.])

Look at the shapes in that sum. origin is (3,) and directions is (2, 3, 3). PyTorch broadcasts the smaller tensor. It reuses the same three origin coordinates for every ray, so you do not have to write a loop or copy the origin six times. Read the shapes before you run the code. When shapes do not line up, the error message points straight at the bug.

To build the full renderer, work through the 0.1 Ray Tracing notebook.

Name the axes

Einstein notation gives tensor axes short names. Once the axes have names, one pattern string can describe a transpose, a sum, or a matrix multiplication. You stop memorizing a separate function for each operation and instead write which axes go in and which axes come out.

Start with the tiny image from the warm-up and name its two axes h and w. The pattern "hw->wh" says the image goes in with height first and comes out with width first. That is a transpose.

The -> arrow separates input axes from output axes. Two rules cover every pattern:

1. An axis named after the -> survives into the result, in the order written there.
2. An axis missing after the -> is summed away.

So "hw->h" sums along width and leaves one total per row, and "hw->" sums away both axes and leaves one number.

import torch

image = torch.tensor([[0, 64, 128], [192, 224, 255]], dtype=torch.float32)
print(torch.einsum("hw->wh", image).shape)  # transpose
print(torch.einsum("hw->h", image))         # sum along width
print(torch.einsum("hw->", image))          # sum everything
# torch.Size([3, 2])
# tensor([192., 671.])
# tensor(863.)

With two inputs there is one more rule. An axis named in both inputs lines up matching values for multiplication. The pattern "hw,w->h" uses every rule at once. w is named in both inputs, so each image value multiplies its matching weight. w is missing after the ->, so the products sum. h survives, so a (2, 3) image and (3,) weights give (2,) row scores.

This weighted sum along one axis is the same pattern the Linear layer uses later on this page.

For the ray grid, use h for height, w for width, and x for the three coordinates in each direction vector. On directions with shape (h, w, x), the h and w axes pick a ray. The x axis stores the three direction components.

Write it with einsum

einsum is Einstein notation you can run. The pattern string lists the axis names for each input, separated by commas, then the output axes after the ->.

Read "hwx,->hwx" this way. The first input has axes h, w, and x. The second input is the scalar t, which has no axes, so its slot before the -> is empty. Every named axis survives to the output, so nothing is summed, and each direction component is multiplied by t.

import torch

directions = torch.tensor([
    [[1, -1, -1], [1, -1, 0], [1, -1, 1]],
    [[1,  1, -1], [1,  1,  0], [1,  1,  1]],
], dtype=torch.float32)
t = torch.tensor(2.0)
scaled = torch.einsum("hwx,->hwx", directions, t)
print(scaled.shape, scaled[0, 1])
# torch.Size([2, 3, 3]) tensor([ 2., -2.,  0.])

ARENA also uses einops.einsum, where the same pattern can use full names: einsum(directions, t, "height width x, -> height width x").

Redo the ray step with einsum

The broadcast version wrote points = origin + 2.0 * directions. The einsum version names the scale step, then adds the origin.

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)
t = torch.tensor(2.0)
scaled = torch.einsum("hwx,->hwx", directions, t)
points = origin + scaled
print(points.shape, points[0, 1])
# torch.Size([2, 3, 3]) tensor([ 2., -2.,  0.])

You should see the same shape and the same point as the broadcast cell. Einsum did not change the computation, only how the pattern is written down. The Linear layer in the next section uses the same reading rule with different axis names.

Follow one digit through the model

Meet MNIST

MNIST contains 70,000 handwritten digits. Each digit is a 28 by 28 grayscale image labeled from 0 to 9.

After the notebook converts an image to a float tensor, each pixel is a number between 0 for black and 1 for white. A human sees the shape of a 3 at once. A neural network learns that shape from 784 separate pixel values.

The Week 1 model

You train a small neural network to recognize handwritten digits from MNIST.

The model receives one image with this shape:

import torch

image = torch.zeros((1, 1, 28, 28))
print(image.shape)
print(image.shape == (1, 1, 28, 28))
# torch.Size([1, 1, 28, 28])
# True

The model runs this computation:

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.

Name the MNIST axes

An MNIST input adds batch and channel axes on top of height and width:

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

The Shape tells you the size and meaning of each axis. The first 1 means one image. The second 1 means one grayscale channel.

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 only rearranges them. No value is removed, added, or changed.

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:

import torch

pixels = torch.zeros((1, 784))
weight = torch.zeros((100, 784))
bias = torch.zeros((100,))
hidden = pixels @ weight.T + bias

print("pixels:", pixels.shape)
print("weight:", weight.shape)
print("bias:", bias.shape)
print("hidden:", hidden.shape)

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. That is the weighted sum from the einsum diagram, with 784 positions instead of 3.

einsum uses the same reading rule as the ray warm-up, now with the full axis names from ARENA's supplied Linear implementation:

import torch
from einops import einsum

pixels = torch.ones((1, 784))
weight = torch.ones((100, 784))
bias = torch.zeros((100,))

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

print(hidden.shape)
print(hidden[0, 0])
# torch.Size([1, 100])
# tensor(784.)

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.

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

Let ReLU bend the computation

ReLU follows the first linear layer. It replaces each negative value with zero and leaves each positive value unchanged:

def relu(x):
    return max(x, 0.0)

hidden = [-3.0, -0.5, 0.0, 1.5, 4.0]
after_relu = [relu(value) 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.

import torch
from torch import nn

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

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


model = SimpleMLP()
image = torch.zeros((1, 1, 28, 28))
logits = model(image)
print(logits.shape)
# torch.Size([1, 10])

This runnable version uses PyTorch's layers. In the first core exercise, you implement the same structure with the Flatten, Linear, and ReLU classes that ARENA supplies.

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.

Run the five-step training loop

Send the loss backward

The forward pass pushes one digit through Flatten, two linear layers, ReLU, and cross entropy to produce a single loss value. While it does this, PyTorch records every operation in a Computational graph . The graph stores which tensors produced each result.

The forward pass builds the graph from left to right. Backpropagation follows it from right to left.

When you call loss.backward(), PyTorch follows the graph from the loss back to every Weight and Bias and computes a Gradient for each one.

A gradient answers one local question. If this parameter increased by a very small amount, how would the loss change? The optimizer uses that answer in step 4, the parameter update.

You do not need to implement backpropagation yourself this week. The 0.4 Backpropagation notebook is stretch work.

Repeat five operations to learn

One training step has five operations:

Scroll the diagram sideways on a small screen.

import torch
from torch import nn
from torch.nn import functional as F

torch.manual_seed(0)
model = nn.Sequential(
    nn.Flatten(),
    nn.Linear(28 * 28, 100),
    nn.ReLU(),
    nn.Linear(100, 10),
)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
images = torch.rand((4, 1, 28, 28))
labels = torch.tensor([0, 1, 2, 3])

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

print("logits shape:", logits.shape)
print("loss:", round(loss.item(), 4))

This check uses four random images. The notebook supplies real MNIST images and labels.

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

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

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

Every optimizer in the course is a variation on one update rule:

θ ← θ − η · g

θ is the current parameters. g is the gradient from loss.backward(). η is the Learning rate , the step size.

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

Optimizer types

Step 4 is not one fixed algorithm. The update rule changes how the model moves across the loss surface, the landscape of loss values over all possible parameters.

The name colors match the balls: SGD is cyan, Momentum is magenta, RMSProp is green, and Adam is blue. The white ball is AdaGrad. Same loss surface, different paths. SGD zig-zags. Momentum overshoots. Adaptive methods equalize steep and flat directions. No optimizer wins everywhere. Learning rate still matters.

The runnable cell above uses SGD. That is enough for the notebook this week. 0.3 Optimization is where you write these update rules and sweep the settings.

Check progress on held-out images

Training loss tells you what happened on the batches used to update the model. You also need images you never train on.

A validation loop runs the model on the test set, the images held back from training, and follows four steps:

1. Run the model inside torch.inference_mode() because validation needs no gradients.
2. Take argmax, the index of the largest of 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, one full pass through the training set.

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

Look ahead to image-aware models

SimpleMLP flattens the image before it learns. Once the pixels are in a line, the model no longer knows which pixels were next to each other.

A Convolution keeps the image axes intact. It applies the same small set of weights at every position in the image.

The training loop you use this week is the same one you will use for convolutions and ResNets. The shapes change. The five training steps stay the same.

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 to implement optimizers and sweep hyperparameters.
6. Use 0.4 Backpropagation to build a small automatic differentiation system.
7. Use 0.1 Ray Tracing for the full renderer.
8. Complete 0.5 VAEs & GANs after the other Chapter 0 work.