Orientation and a first digit classifier

This week is our foundation, and we are going from tensors all the way to backpropagation, building the core mental model that supports everything else in the course.

The reason we start here is that tensors are the basic data structure underneath almost everything we will do, including images, activations, gradients, and the data flowing through models. If you get comfortable reading shapes and understanding what each axis means, the rest of the course becomes much easier.

Over the next sections, we will move from tensor basics, to some PyTorch practice with rays, to building and training a simple model, and then into optimization and backpropagation. The goal is not to rush ahead, but to build intuition step by step so these ideas feel solid when we see them again later.

Warm up with tensors and a simple ray example

What is a tensor?

A tensor is a multidimensional array of numbers. Its shape lists the size of each axis and helps you identify the data it holds.

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, you could use one loop over the rows and another over the columns. PyTorch applies one operation to every selected value at once. This is a vectorized operation. It avoids one Python step per pixel and is faster on large tensors.

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 value t controls how far the point moves along D:

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. Check the shapes before you run the code. When an operation fails, the error message often shows which shapes do not match.

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

In "hwx,->hwx", 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 remains in the output. The operation multiplies each direction component by t without summing any axis.

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 point as in the broadcasting cell. Einsum expresses the same computation with named axes. The Linear layer in the next section uses the same axis rules with different 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 has a value from 0 to 1. Here, 0 is black and 1 is white. The model receives 784 separate pixel values and learns which patterns identify each digit.

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 produces ten logits. Each logit 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 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.

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 first 1 means one image in the batch. 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 axis rules 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.)

The pattern follows three rules:

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

Apply ReLU between the linear layers

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 have the same effect as one linear layer. ReLU changes the values 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.

Choose the largest logit

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.

import torch

logits = torch.tensor([[0.2, -0.4, 1.1, 0.6, -0.8, 3.2, 0.0, 0.9, -0.2, 0.3]])
prediction = logits.argmax().item()

print(prediction)
# 5

argmax returns the index of the largest logit. The largest value is 3.2 at index 5, so the predicted digit is 5.

Softmax turns logits into probabilities that add to 1. During training, pass the raw logits to cross-entropy. F.cross_entropy computes the loss directly from those logits, 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 example uses four random images. The notebook supplies real MNIST images and labels.

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

θ ← θ − η · g

θ represents 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

Different optimizers use the gradient in different ways. A loss surface plots parameter values against the resulting loss.

The name colors match the balls: SGD is cyan, Momentum is magenta, RMSProp is green, and Adam is blue. AdaGrad is white. Each optimizer follows a different path across the same loss surface. SGD zigzags, while Momentum overshoots. RMSProp and Adam adjust their steps across steep and flat directions. No optimizer wins everywhere. The learning rate still affects every path.

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 how the model performed on the batches used to update it. You also need images that were held back from training.

The notebook uses its test set for validation. The validation loop 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.

A ResNet stacks residual blocks, groups of convolutional layers with a direct path around them. In the simplified block below, two 3 by 3 convolutions with a ReLU between them compute an update F(x) with the same shape as x. A skip connection, a direct path around the convolutions, carries x to the end. The block adds x and F(x), then applies ReLU:

output = ReLU(x + F(x))

The skip connection gives information and gradients a direct path through the block. This makes deeper networks easier to train.

The full block in the notebook also applies batch normalization after each convolution to normalize and rescale each channel during training.

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 at least 95 percent test accuracy 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.