PT02: Logistic regression in PyTorch (instructor version)

Given a few measurements of a penguin, can we predict its species? In this lab we build a linear classifier, train it with gradient descent and evaluate it.

You should be comfortable with PT01 up to its short check: tensor operations, backward() and parameter updates. Data loading and plotting are provided. There are six activities. In the first three you write the forward pass, the loss and the training step. Then you investigate why the largest learning rate wins on this data, diagnose a faulty training step, and check when a mini-batch update equals a full-batch one. Along the way we switch to PyTorch's standard components and train on mini-batches. After the MLP lecture, Section 9 replaces the linear model with an MLP, keeping the rest of the code.

Two optional sections close the notebook. The first uses per-example gradients to look inside the trained classifier. The second is an exercise for after the MLP lecture, where we reuse the same training loop with an MLP to fit an image.

This is the instructor version. Reference implementations are included, and unfinished activities fall back to them, so the whole notebook runs. The student version does not contain these cells.

Try to solve these exercises without AI coding assistants.

As in PT01, jaxtyping annotations document the shapes of function inputs and outputs, for example Float[torch.Tensor, "batch features"]. Nothing checks them at runtime, and they say nothing about intermediate shapes: those you have to work out yourself.

If the import fails, run %pip install jaxtyping in a new cell, then rerun the imports.

from pathlib import Path
import copy
import hashlib
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch
from jaxtyping import Float, Int
from typing import Optional
import torch.nn.functional as F
from sklearn.model_selection import train_test_split

torch.manual_seed(7)
torch.set_printoptions(precision=3, sci_mode=False)
print("PyTorch", torch.__version__)

Matplotlib is building the font cache; this may take a moment.

PyTorch 2.14.1

1. The data

We use the Palmer Penguins dataset: four measurements (bill length, bill depth, flipper length and body mass) for penguins of three species, Adelie, Chinstrap and Gentoo. Each row is one penguin. We drop the two rows with missing measurements. The next cell reads the copy in data/ if it exists, and downloads it otherwise.

DATA_URL = "https://raw.githubusercontent.com/allisonhorst/palmerpenguins/main/inst/extdata/penguins.csv"
DATA_SHA256 = "f204db2c753b0937caac3cb35258562c14f073e4bbc76be24b4c51ce22767a93"
data_path = Path("data/penguins.csv")
if not data_path.exists():
    from urllib.request import urlopen
    data_path.parent.mkdir(parents=True, exist_ok=True)
    data_path.write_bytes(urlopen(DATA_URL, timeout=30).read())
assert hashlib.sha256(data_path.read_bytes()).hexdigest() == DATA_SHA256, "Dataset changed; use the supplied data/penguins.csv."

features = ["bill_length_mm", "bill_depth_mm", "flipper_length_mm", "body_mass_g"]
classes = ["Adelie", "Chinstrap", "Gentoo"]
penguins = pd.read_csv(data_path).dropna(subset=features + ["species"])
X = torch.tensor(penguins[features].to_numpy(), dtype=torch.float32)
y = torch.tensor(pd.Categorical(penguins.species, categories=classes).codes.copy(), dtype=torch.long)
assert X.shape == (342, 4) and y.shape == (342,)
penguins[features + ["species"]].head()

	bill_length_mm	bill_depth_mm	flipper_length_mm	body_mass_g	species
0	39.1	18.7	181.0	3750.0	Adelie
1	39.5	17.4	186.0	3800.0	Adelie
2	40.3	18.0	195.0	3250.0	Adelie
4	36.7	19.3	193.0	3450.0	Adelie
5	39.3	20.6	190.0	3650.0	Adelie

We split the data into a training set (60%), a validation set (about 20%) and a test set (the remaining 20%). The split is stratified: each set has roughly the same class proportions.

The features have very different scales (body mass is in grams, the others in millimeters), so we standardize them. The mean and standard deviation come from the training set only, and are then applied to all three sets. This way the validation and test examples do not influence anything we compute during training.

indices = np.arange(len(y))
train_idx, rest_idx = train_test_split(indices, test_size=0.4, stratify=y.numpy(), random_state=7)
val_idx, test_idx = train_test_split(rest_idx, test_size=0.5, stratify=y[rest_idx].numpy(), random_state=7)

Xtrain_raw, ytrain = X[train_idx], y[train_idx]
Xval_raw, yval = X[val_idx], y[val_idx]
Xtest_raw, ytest = X[test_idx], y[test_idx]
mean = Xtrain_raw.mean(dim=0)
scale = Xtrain_raw.std(dim=0, correction=0).clamp_min(1e-8)
Xtrain = (Xtrain_raw - mean) / scale
Xval = (Xval_raw - mean) / scale
Xtest = (Xtest_raw - mean) / scale

print("Examples:", len(ytrain), "training;", len(yval), "validation;", len(ytest), "test")
print("Training means:", Xtrain.mean(dim=0))
torch.testing.assert_close(Xtrain.mean(dim=0), torch.zeros(4), atol=2e-6, rtol=0)

Examples: 205 training; 68 validation; 69 test
Training means: tensor([-0.000, -0.000,  0.000, -0.000])

fig, ax = plt.subplots(figsize=(6, 4))
for c, name in enumerate(classes):
    points = Xtrain_raw[ytrain == c]
    ax.scatter(points[:, 0], points[:, 1], label=name, marker=["o", "^", "s"][c], s=24, alpha=0.8)
ax.set(xlabel="Bill length (mm)", ylabel="Bill depth (mm)")
ax.legend()
fig.tight_layout()
plt.show()

2. The model

Recall the batched linear model from PT01,

$$Z=XW^\top+b,$$

with $X\sim(B,D)$, $W\sim(C,D)$, $b\sim(C,)$ and $Z\sim(B,C)$. Here $D=4$ features and $C=3$ classes. The outputs are logits; we apply a softmax only when we need probabilities.

Activity 1: a batched forward pass

Implement linear_logits with tensor operations and run the check on the small example. How is the bias added to each example? The model class that uses your function is provided after the check.

def linear_logits(X: Float[torch.Tensor, "batch features"], W: Float[torch.Tensor, "classes features"], b: Float[torch.Tensor, "classes"]) -> Optional[Float[torch.Tensor, "batch classes"]]:
    # Replace None with your batched computation.
    return None

# Reference used only when linear_logits returns None.
def reference_linear_logits(X: Float[torch.Tensor, "batch features"], W: Float[torch.Tensor, "classes features"], b: Float[torch.Tensor, "classes"]) -> Float[torch.Tensor, "batch classes"]:
    return X @ W.T + b

def forward_logits(X: Float[torch.Tensor, "batch features"], W: Float[torch.Tensor, "classes features"], b: Float[torch.Tensor, "classes"]) -> Float[torch.Tensor, "batch classes"]:
    result = linear_logits(X, W, b)
    return reference_linear_logits(X, W, b) if result is None else result

small_X = torch.tensor([[1., 2.], [3., 4.]])
small_W = torch.tensor([[1., 0.], [0., 1.], [-1., 1.]])
small_b = torch.tensor([0.5, -0.5, 1.])
expected = torch.tensor([[1.5, 1.5, 2.], [3.5, 3.5, 2.]])
torch.testing.assert_close(forward_logits(small_X, small_W, small_b), expected)
print("Using", "reference" if linear_logits(small_X, small_W, small_b) is None else "your", "forward pass")

Using reference forward pass

A torch.nn.Module keeps track of the parameters of a model. Wrapping a tensor in nn.Parameter registers it as a parameter and sets requires_grad=True. We initialize everything to zero. This is fine for a linear classifier, but it would not work for the hidden layers of an MLP, as Section 9 shows.

class LogisticRegression(torch.nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.weight = torch.nn.Parameter(torch.zeros(output_dim, input_dim))
        self.bias = torch.nn.Parameter(torch.zeros(output_dim))

    def forward(self, X: Float[torch.Tensor, "batch features"]) -> Float[torch.Tensor, "batch classes"]:
        return forward_logits(X, self.weight, self.bias)

model = LogisticRegression(Xtrain.shape[1], len(classes))
assert model(Xtrain[:5]).shape == (5, 3)
print("Parameter tensors:", len(list(model.parameters())))
print("Trainable scalar parameters:", sum(p.numel() for p in model.parameters()))

Parameter tensors: 2
Trainable scalar parameters: 15

3. Loss and accuracy

For a vector of logits $z$ and correct class $y$, the cross-entropy is

$$\ell(z,y)=\log\sum_c\exp(z_c)-z_y.$$

Activity 2: the loss

Implement student_cross_entropy from this formula, averaged over the examples in the batch. Use torch.logsumexp for the first term, and select the logit of the correct class from each row. Do not call F.cross_entropy: we use it below to check your result.

Before writing code, work out the shapes of the first term, of the selected logits and of the final loss. The accuracy function is provided.

def student_cross_entropy(logits: Float[torch.Tensor, "batch classes"], targets: Int[torch.Tensor, "batch"]) -> Optional[Float[torch.Tensor, ""]]:
    # Replace None with the batched loss.
    return None

# Reference used only while the loss activity is unfilled.
def reference_cross_entropy(logits: Float[torch.Tensor, "batch classes"], targets: Int[torch.Tensor, "batch"]) -> Float[torch.Tensor, ""]:
    rows = torch.arange(len(targets), device=targets.device)
    return (torch.logsumexp(logits, dim=1) - logits[rows, targets]).mean()

def cross_entropy(logits: Float[torch.Tensor, "batch classes"], targets: Int[torch.Tensor, "batch"]) -> Float[torch.Tensor, ""]:
    result = student_cross_entropy(logits, targets)
    return reference_cross_entropy(logits, targets) if result is None else result

def accuracy(logits: Float[torch.Tensor, "batch classes"], targets: Int[torch.Tensor, "batch"]) -> Float[torch.Tensor, ""]:
    return (logits.argmax(dim=1) == targets).float().mean()

logits = model(Xtrain)
print("Initial loss:", cross_entropy(logits, ytrain).item())
print("Initial accuracy:", accuracy(logits, ytrain).item())
# Zero logits give uniform probabilities and loss log(C).
torch.testing.assert_close(cross_entropy(logits, ytrain), torch.tensor(float(np.log(3))))
print("Using", "reference" if student_cross_entropy(logits, ytrain) is None else "your", "loss")

Initial loss: 1.0986120700836182
Initial accuracy: 0.4390243887901306
Using reference loss

At initialization all logits are zero, so the predicted probabilities are uniform and the loss is $\log 3$. The accuracy is not one third, though: argmax breaks ties by picking the first class (Adelie), so the initial accuracy is the fraction of Adelie penguins in the training set.

The next cell compares your loss and its gradient with PyTorch's. The second test uses very large logits, where computing log(softmax(...)) directly would give infinite values.

# Check unequal losses as well as extreme logits: averaging must be over examples.
for values, labels in [([[2., -1., 0.], [0., 1., -2.]], [0, 2]),
                       ([[1000., 0., -1000.], [-1000., 0., 1000.]], [2, 0])]:
    probe = torch.tensor(values, requires_grad=True)
    targets = torch.tensor(labels)
    ours = cross_entropy(probe, targets)
    builtin = F.cross_entropy(probe, targets)
    torch.testing.assert_close(ours, builtin)
    g_ours, = torch.autograd.grad(ours, probe)
    g_builtin, = torch.autograd.grad(builtin, probe)
    torch.testing.assert_close(g_ours, g_builtin)
    assert torch.isfinite(ours)
print("Loss values and gradients agree.")

Loss values and gradients agree.

4. One training step

One step of gradient descent has four parts, in this order:

1. clear the gradients left by the previous step (set the .grad of each parameter to None);
2. compute the loss on the data;
3. call backward();
4. subtract the learning rate times its gradient from each parameter.

Only the update runs inside torch.no_grad(). You saw all four parts in PT01; here you put them together.

# Instructor demonstration of one full step; not in the student version.
model = LogisticRegression(Xtrain.shape[1], len(classes))
for parameter in model.parameters():
    parameter.grad = None
loss = cross_entropy(model(Xtrain), ytrain)
loss.backward()
print("Weight gradient shape:", model.weight.grad.shape)
print("Bias gradient shape:", model.bias.grad.shape)
with torch.no_grad():
    for parameter in model.parameters():
        parameter -= 0.1 * parameter.grad
print("Loss before:", loss.item())
print("Loss after:", cross_entropy(model(Xtrain), ytrain).item())

Weight gradient shape: torch.Size([3, 4])
Bias gradient shape: torch.Size([3])
Loss before: 1.0986120700836182
Loss after: 0.9933529496192932

Activity 3: train the classifier

Write training_step. It makes one step on the data it receives, changes the model in place and returns the loss computed before the update, as a Python number. Loop over model.parameters(), so that the same code updates the weight and the bias. The data are arguments rather than global variables, so later we can call the same function on mini-batches.

The next cell checks your function over two consecutive steps; the second step catches gradients that were not cleared. Then run the training loop and look at the training and validation curves together. Leave the test set alone until the end.

def training_step(model: torch.nn.Module, X: Float[torch.Tensor, "batch features"], y: Int[torch.Tensor, "batch"], lr: float) -> Optional[float]:
    # Replace None with your training step.
    return None

def reference_training_step(model: torch.nn.Module, X: Float[torch.Tensor, "batch features"], y: Int[torch.Tensor, "batch"], lr: float) -> float:
    for parameter in model.parameters():
        parameter.grad = None
    loss = cross_entropy(model(X), y)
    loss.backward()
    with torch.no_grad():
        for parameter in model.parameters():
            parameter -= lr * parameter.grad
    return loss.item()

def run_step(model: torch.nn.Module, X: Float[torch.Tensor, "batch features"], y: Int[torch.Tensor, "batch"], lr: float) -> float:
    # A skipped exercise returns None without modifying the model.
    before = [p.detach().clone() for p in model.parameters()]
    result = training_step(model, X, y, lr)
    if result is None:
        assert all(torch.equal(p, saved) for p, saved in zip(model.parameters(), before)), "Your step changed parameters but returned None. Return loss.item()."
        # Clear any gradients left by a partially completed exercise.
        return reference_training_step(model, X, y, lr)
    return result

@torch.no_grad()
def evaluate(model, X, y):
    model.eval()
    logits = model(X)
    return cross_entropy(logits, y).item(), accuracy(logits, y).item()

# Compare two consecutive updates: the second detects uncleared gradients.
check_model = LogisticRegression(4, 3)
reference_model = copy.deepcopy(check_model)
for _ in range(2):
    actual = run_step(check_model, Xtrain[:16], ytrain[:16], 0.1)
    expected_loss = reference_training_step(reference_model, Xtrain[:16], ytrain[:16], 0.1)
    assert isinstance(actual, float), "Return loss.item(), not a tensor."
    assert abs(actual - expected_loss) < 1e-6
    for parameter, expected_parameter in zip(check_model.parameters(), reference_model.parameters()):
        torch.testing.assert_close(parameter, expected_parameter)
print("Both updates agree with the reference.")

Both updates agree with the reference.

def fit(model, Xtrain, ytrain, Xval, yval, lr=0.1, epochs=200):
    history = {name: [] for name in ["train_loss", "val_loss", "train_accuracy", "val_accuracy"]}
    for _ in range(epochs):
        model.train()
        run_step(model, Xtrain, ytrain, lr)
        # Both curves measure the model after the same update.
        train_loss, train_acc = evaluate(model, Xtrain, ytrain)
        val_loss, val_acc = evaluate(model, Xval, yval)
        for key, value in zip(history, [train_loss, val_loss, train_acc, val_acc]):
            history[key].append(value)
    return history

model = LogisticRegression(4, 3)
history = fit(model, Xtrain, ytrain, Xval, yval)
fig, axes = plt.subplots(1, 2, figsize=(10, 3.5))
epochs = range(1, len(history["train_loss"]) + 1)
for split, style in [("train", "-"), ("val", "--")]:
    axes[0].plot(epochs, history[f"{split}_loss"], style, label=split)
    axes[1].plot(epochs, history[f"{split}_accuracy"], style, label=split)
axes[0].set(xlabel="Epoch", ylabel="Cross-entropy")
axes[1].set(xlabel="Epoch", ylabel="Accuracy", ylim=(0, 1.03))
for ax in axes:
    ax.legend()
fig.tight_layout()
plt.show()
print("Final validation loss and accuracy:", evaluate(model, Xval, yval))

Final validation loss and accuracy: (0.12177293747663498, 0.9852941036224365)

Each update uses the whole training set, so here one update is one epoch (one pass through the data). With mini-batches, an epoch contains several updates.

The loss can go down while the accuracy stays flat: the probability of the correct class can grow without changing which class has the largest logit. Can you find such an interval in your curves?

Comparing learning rates

Every run starts from the same parameters and makes 200 full-batch updates; only the learning rate changes. Before running the cell, predict how the curves for 0.001, 0.1 and 10.0 will differ. Does accuracy tell the same story as the loss?

We keep the run with the lowest final validation loss. The test set plays no part in this choice.

learning_rates = [0.001, 0.1, 10.0]
initial_model = LogisticRegression(4, 3)
lr_models, lr_histories = {}, {}
fig, axes = plt.subplots(1, 2, figsize=(10, 3.5))
for rate, style in zip(learning_rates, ["-", "--", ":"]):
    candidate = copy.deepcopy(initial_model)
    trace = fit(candidate, Xtrain, ytrain, Xval, yval, lr=rate, epochs=200)
    lr_models[rate], lr_histories[rate] = candidate, trace
    for ax, split in zip(axes, ["train", "val"]):
        ax.plot(range(1, 201), trace[f"{split}_loss"], style, label=f"lr={rate:g}")
    print(f"lr={rate:g}: validation loss={trace['val_loss'][-1]:.4f}, accuracy={trace['val_accuracy'][-1]:.3f}")
for ax, title in zip(axes, ["Training", "Validation"]):
    ax.set(xlabel="Update", ylabel="Cross-entropy", title=title, yscale="log")
    ax.legend()
fig.tight_layout()
plt.show()
selected_lr = min(learning_rates, key=lambda rate: lr_histories[rate]["val_loss"][-1])
model = lr_models[selected_lr]
print("Selected learning rate:", selected_lr)

lr=0.001: validation loss=0.9241, accuracy=0.868

lr=0.1: validation loss=0.1218, accuracy=0.985

lr=10: validation loss=0.0030, accuracy=1.000

Selected learning rate: 10.0

On this dataset the largest rate wins clearly. Do not read this as a general rule: the next activity asks why it happens here.

Activity 4: why does the largest rate win?

Write down a prediction before each step, then check it.

1. Which of the three runs has the largest weight norm? Print lr_models[rate].weight.norm() and the training accuracy of each run.
2. Copy the lr=10 model with copy.deepcopy and keep training it with fit for 2000 more updates, recording the weight norm every 200 updates. Does the training loss reach zero? Does the norm level off?
3. The norm keeps changing, but does the direction W / ‖W‖ keep changing too? Compute the cosine between the directions at consecutive checkpoints.
4. Train from initial_model with lr=100 and lr=1000, and print the first ten training losses. Does anything break? In Section 9 we will see what a rate of 10 does to an MLP.

Do not overwrite model or lr_models: Section 8 evaluates the selected run.

# Activity 4: your experiments here.

# Reference for Activity 4.
for rate in learning_rates:
    print(f"lr={rate:g}: weight norm {lr_models[rate].weight.norm().item():.2f}, "
          f"training accuracy {evaluate(lr_models[rate], Xtrain, ytrain)[1]:.3f}")

longer = copy.deepcopy(lr_models[10.0])
norms, directions, train_losses = [], [], []
for _ in range(10):
    trace = fit(longer, Xtrain, ytrain, Xval, yval, lr=10.0, epochs=200)
    norms.append(longer.weight.norm().item())
    directions.append((longer.weight / longer.weight.norm()).detach().flatten())
    train_losses.append(trace["train_loss"][-1])
cosines = [float(a @ b) for a, b in zip(directions, directions[1:])]
print("Norm every 200 updates:", [round(n, 1) for n in norms])
print("Training loss:", [f"{v:.1e}" for v in train_losses])
print("Cosine between consecutive directions:", [round(c, 4) for c in cosines])
assert all(b > a for a, b in zip(norms, norms[1:]))
assert all(b < a for a, b in zip(train_losses, train_losses[1:])) and train_losses[-1] > 0
assert cosines[-1] > 0.999

for rate in [100.0, 1000.0]:
    trial = copy.deepcopy(initial_model)
    trace = fit(trial, Xtrain, ytrain, Xval, yval, lr=rate, epochs=200)
    print(f"lr={rate:g}: first losses", [round(v, 2) for v in trace["train_loss"][:10]],
          f"final {trace['train_loss'][-1]:.1e}")
    assert np.isfinite(trace["train_loss"]).all()

lr=0.001: weight norm 0.19, training accuracy 0.932
lr=0.1: weight norm 2.99, training accuracy 0.980
lr=10: weight norm 13.83, training accuracy 1.000

Norm every 200 updates: [15.8, 17.2, 18.4, 19.4, 20.4, 21.2, 21.9, 22.6, 23.3, 23.8]
Training loss: ['7.6e-03', '6.2e-03', '5.3e-03', '4.7e-03', '4.3e-03', '3.9e-03', '3.6e-03', '3.3e-03', '3.1e-03', '3.0e-03']
Cosine between consecutive directions: [0.9986, 0.9995, 0.9998, 0.9999, 0.9999, 1.0, 1.0, 1.0, 1.0]
lr=100: first losses [3.47, 1.15, 0.45, 0.27, 0.15, 0.08, 0.03, 0.01, 0.01, 0.0] final 3.1e-04
lr=1000: first losses [34.63, 11.82, 4.65, 2.78, 1.47, 0.75, 0.21, 0.02, 0.0, 0.0] final 2.1e-06

Discussion (open after the activity)

With these four features the three species are linearly separable on the training set: the lr=10 run reaches 100% training accuracy. Once a linear classifier separates the data, scaling up its weights makes every correct logit larger relative to the others, so the cross-entropy keeps falling as the norm grows. The loss never reaches zero and the norm never levels off, though both change more and more slowly. The direction, instead, settles: training keeps pushing the same separator further out. A larger step simply travels this path faster, which is why lr=10 wins.

Even rates of 100 and 1000 do not break anything. The first update overshoots and the loss jumps, but the loss is convex in the parameters and the data are separable, so training recovers. This is a property of this problem, not of gradient descent: with an MLP, or with classes that overlap, the same rates behave very differently.

5. Investigating a training run

Activity 5: a different training step

A colleague gives you the function below. It runs, and at first its loss may even fall faster than yours. Is it the same gradient descent?

Compare it with your training_step, starting from the same parameters and data. Design a small experiment that tells the two apart: write down what you expect, run it, then fix the function. A decreasing loss alone does not settle the question. The cell below only defines the function; the experiment is up to you.

def suspect_step(model: torch.nn.Module, X: Float[torch.Tensor, "batch features"], y: Int[torch.Tensor, "batch"], lr: float) -> float:
    loss = cross_entropy(model(X), y)
    loss.backward()
    with torch.no_grad():
        for parameter in model.parameters():
            parameter -= lr * parameter.grad
    return loss.item()

# Activity 5: your experiment here. What do you expect, and what do you observe?

# Reference probe: differentiate twice without updating the parameters.
probe_model = LogisticRegression(4, 3)
cross_entropy(probe_model(Xtrain), ytrain).backward()
first = probe_model.weight.grad.clone()
cross_entropy(probe_model(Xtrain), ytrain).backward()
torch.testing.assert_close(probe_model.weight.grad, 2 * first)
print("Same parameters, same data, twice the stored gradient.")
# Clear gradients at the start of suspect_step, as in reference_training_step.

Same parameters, same data, twice the stored gradient.

# Reference comparison to show after the investigation.
normal = LogisticRegression(4, 3)
suspect = copy.deepcopy(normal)
normal_losses, suspect_losses = [], []
normal_norms, suspect_norms = [], []
for _ in range(30):
    run_step(normal, Xtrain, ytrain, 0.1)
    suspect_step(suspect, Xtrain, ytrain, 0.1)
    normal_losses.append(evaluate(normal, Xtrain, ytrain)[0])
    suspect_losses.append(evaluate(suspect, Xtrain, ytrain)[0])
    normal_norms.append(normal.weight.grad.norm().item())
    suspect_norms.append(suspect.weight.grad.norm().item())
fig, axes = plt.subplots(1, 2, figsize=(10, 3.5))
updates = range(1, len(normal_losses) + 1)
axes[0].plot(updates, normal_losses, label="Working step")
axes[0].plot(updates, suspect_losses, "--", label="Suspect step")
axes[0].set(xlabel="Update", ylabel="Training cross-entropy")
axes[1].plot(updates, normal_norms, label="Working step")
axes[1].plot(updates, suspect_norms, "--", label="Suspect step")
axes[1].set(xlabel="Update", ylabel="Norm of weight.grad")
for ax in axes:
    ax.legend()
fig.tight_layout()
plt.show()

Discussion (open after the investigation)

suspect_step never clears the gradients, so each update uses the sum of all gradients computed so far, most of them at older parameter values. This is no longer gradient descent. The loss can still go down, which is why the loss curve alone does not reveal the bug.

Accumulating gradients is sometimes what we want, for example to combine several mini-batches into one update, as long as the loss is scaled correctly (see case 3 of the autograd detective in PT01). Here, instead, we update after every backward() and never clear.

6. Using the standard components

In practice we do not write these pieces ourselves: nn.Linear holds the weight and bias, F.cross_entropy computes the loss, and an optimizer such as torch.optim.SGD performs the update. They should compute exactly what our code computes.

To check it, we copy the parameters of our model into the standard one and compare a single update. If the two models started from different random parameters, we could not tell implementation differences apart from initialization differences.

manual = LogisticRegression(4, 3)
standard = torch.nn.Linear(4, 3)
with torch.no_grad():
    standard.weight.copy_(manual.weight)
    standard.bias.copy_(manual.bias)
optimizer = torch.optim.SGD(standard.parameters(), lr=0.1)

manual_loss = run_step(manual, Xtrain, ytrain, 0.1)
optimizer.zero_grad(set_to_none=True)
standard_loss = F.cross_entropy(standard(Xtrain), ytrain)
standard_loss.backward()
optimizer.step()

torch.testing.assert_close(standard.weight, manual.weight)
torch.testing.assert_close(standard.bias, manual.bias)
torch.testing.assert_close(standard(Xval), manual(Xval))
print("Manual and standard updates agree.")

Manual and standard updates agree.

7. From full batches to mini-batches

A TensorDataset pairs each row of features with its label. A DataLoader splits the dataset into batches and, with shuffle=True, uses a new order in every epoch.

Before running: with 205 training examples and batches of 32, how many updates are there in one epoch? How large is the last batch?

We train a new model from the same initial parameters with lr=0.1 and compare it with the full-batch run at the same rate. After the same number of epochs, both models have seen the data equally often, but the mini-batch model has made many more updates.

from torch.utils.data import TensorDataset, DataLoader

train_loader = DataLoader(TensorDataset(Xtrain, ytrain), batch_size=32,
                          shuffle=True, drop_last=False,
                          generator=torch.Generator().manual_seed(7))
mini_model = copy.deepcopy(initial_model)
mini_history = {"train_loss": [], "val_loss": []}
mini_updates = 0
batch_sizes = []
for epoch in range(30):
    mini_model.train()
    for batch_X, batch_y in train_loader:
        run_step(mini_model, batch_X, batch_y, lr=0.1)
        mini_updates += 1
        if epoch == 0:
            batch_sizes.append(len(batch_y))
    # Evaluate the current model on each complete split. Do not average
    # unweighted batch losses: the final batch contains fewer examples.
    mini_history["train_loss"].append(evaluate(mini_model, Xtrain, ytrain)[0])
    mini_history["val_loss"].append(evaluate(mini_model, Xval, yval)[0])
print("Batch sizes in one epoch:", batch_sizes)
print("Epochs: 30; updates:", mini_updates)
fig, ax = plt.subplots(figsize=(6, 4))
ax.plot(range(1, 31), mini_history["val_loss"], label="Mini-batch, 32 examples")
ax.plot(range(1, 31), lr_histories[0.1]["val_loss"][:30], "--", label="Full batch")
ax.set(xlabel="Epoch", ylabel="Validation cross-entropy")
ax.legend()
fig.tight_layout()
plt.show()

Batch sizes in one epoch: [32, 32, 32, 32, 32, 32, 13]
Epochs: 30; updates: 210

Activity 6: a batch as large as the training set

Why do two batches give different gradients for the same parameters? Write down your answer, then test a consequence of it. Build a DataLoader with batch_size=len(ytrain) and shuffle=True, and show that one epoch with it, starting from initial_model, gives the same parameters as one full-batch run_step from the same start. Does the shuffling matter here? What would change with batch_size=len(ytrain) - 1?

# Activity 6: your check here.

# Reference for Activity 6.
full_loader = DataLoader(TensorDataset(Xtrain, ytrain), batch_size=len(ytrain), shuffle=True,
                         generator=torch.Generator().manual_seed(7))
one_epoch, one_step = copy.deepcopy(initial_model), copy.deepcopy(initial_model)
updates_in_epoch = 0
for batch_X, batch_y in full_loader:
    run_step(one_epoch, batch_X, batch_y, lr=0.1)
    updates_in_epoch += 1
run_step(one_step, Xtrain, ytrain, lr=0.1)
assert updates_in_epoch == 1
for loader_parameter, full_parameter in zip(one_epoch.parameters(), one_step.parameters()):
    torch.testing.assert_close(loader_parameter, full_parameter)
print("One shuffled full-size batch gives the full-batch update.")

One shuffled full-size batch gives the full-batch update.

Discussion (open after the activity)

Each batch gradient is the average of the per-example gradients in that batch, so different examples give different averages. With a single batch containing every example, the average is the full-batch gradient, and shuffling only changes the order of the terms in the sum: the result agrees up to floating-point rounding, which is why the check uses assert_close and not exact equality. With batch_size=len(ytrain) - 1 an epoch has two updates, the second computed from a single example, so its gradient is far noisier than the first.

8. Final evaluation

We now evaluate model, the full-batch run selected by validation loss, on the test set. This is the first time we use the test set.

For reference, we compare its accuracy with a trivial baseline that always predicts the most common class in the training set. The confusion matrix shows which species are mistaken for which. If you go back and change the model after seeing these numbers, the test set no longer gives a fair final evaluation.

test_loss, test_accuracy = evaluate(model, Xtest, ytest)
majority_class = torch.bincount(ytrain).argmax()
baseline_accuracy = (ytest == majority_class).float().mean().item()
print(f"Test cross-entropy: {test_loss:.3f}")
print(f"Test accuracy: {test_accuracy:.3f}")
print(f"Majority-class accuracy: {baseline_accuracy:.3f}")

with torch.no_grad():
    predicted = model(Xtest).argmax(dim=1)
confusion = torch.bincount(ytest * 3 + predicted, minlength=9).reshape(3, 3)
confusion_table = pd.DataFrame(confusion.numpy(), index=classes, columns=classes)
confusion_table.index.name = "True species"
confusion_table.columns.name = "Predicted species"
confusion_table

Test cross-entropy: 0.035
Test accuracy: 0.971
Majority-class accuracy: 0.449

Predicted species	Adelie	Chinstrap	Gentoo
True species
Adelie	30	1	0
Chinstrap	1	12	0
Gentoo	0	0	25

Pick one misclassified penguin and look at its logits or probabilities. Was the model confident? What can you conclude from a single example, and what would need more evidence?

The same training loop works unchanged when we replace the linear model with an MLP. Section 9 shows this; come back to it after the MLP lecture.

9. From logistic regression to an MLP

Come back to this section after the MLP lecture. We replace the linear model with an MLP with one hidden layer of 16 ReLU units. Nothing else changes: fit, evaluate, your training_step and your cross_entropy only need a module that maps a batch of features to a batch of logits, and they work as they are.

nn.Sequential chains layers. Unlike our LogisticRegression, each nn.Linear starts from small random weights. We train with the same learning rate and number of updates as the linear run at lr=0.1, and compare the two on the validation set. The test set stays untouched: the model we evaluated in Section 8 remains our final one.

def make_mlp(width: int = 16) -> torch.nn.Module:
    return torch.nn.Sequential(
        torch.nn.Linear(4, width), torch.nn.ReLU(), torch.nn.Linear(width, len(classes))
    )

torch.manual_seed(7)
mlp = make_mlp()
mlp_history = fit(mlp, Xtrain, ytrain, Xval, yval, lr=0.1, epochs=200)
print("MLP parameters:", sum(p.numel() for p in mlp.parameters()))
print(f"Validation loss after 200 updates: MLP {mlp_history['val_loss'][-1]:.4f}, "
      f"linear {lr_histories[0.1]['val_loss'][-1]:.4f}")

fig, ax = plt.subplots(figsize=(6, 4))
ax.plot(range(1, 201), lr_histories[0.1]["val_loss"], label="Linear")
ax.plot(range(1, 201), mlp_history["val_loss"], "--", label="MLP, 16 hidden units")
ax.set(xlabel="Update", ylabel="Validation cross-entropy", yscale="log", title="lr=0.1")
ax.legend()
fig.tight_layout()
plt.show()

MLP parameters: 131
Validation loss after 200 updates: MLP 0.0612, linear 0.1218

At the same learning rate the MLP reaches a lower validation loss, but remember that the linear model at lr=10 was already close to perfect on this data. The point of this section is how little code changed. Two experiments show where the MLP does behave differently:

1. Train a new MLP with lr=10.0, the rate that won the comparison for the linear model. Before running, predict the training curve. What happens, and why could the linear model take such large steps?
2. In Section 2 we initialized the linear model with zeros. Do the same for the MLP: after building it, set every parameter to zero inside torch.no_grad(), then train with lr=0.1. Predict the final loss. Which parameters change during training, and which stay at zero?

# Try the two experiments here.

Section 9: answers

1. With lr=10.0 the MLP's training loss jumps from about 1 to very large values within a few updates. In the recorded run it eventually comes back down, but the curve is erratic and the final validation loss is much worse than at lr=0.1; at lr=30 the loss becomes NaN. The linear model's loss is convex and its gradients stay bounded, so on this separable data large steps only push the weights further in a good direction. The MLP's loss is not convex, and a large step can land in a region with much larger gradients, which makes the next step larger still.
2. With all parameters at zero, every hidden unit outputs zero and receives zero gradient, and so does the output weight. Only the output bias learns, so the model can only learn the class frequencies: the loss stalls near 1.05 and the model predicts one class for every penguin. Random initialization breaks this symmetry. A zero start was fine for logistic regression because there are no hidden units to tell apart.

torch.manual_seed(7)
unstable_mlp = make_mlp()
unstable_history = fit(unstable_mlp, Xtrain, ytrain, Xval, yval, lr=10.0, epochs=200)
unstable_losses = torch.tensor(unstable_history["train_loss"])
print("lr=10: largest training loss", unstable_losses.nan_to_num(float("inf")).max().item())
assert not torch.isfinite(unstable_losses).all() or unstable_losses.max() > 100

zero_mlp = make_mlp()
with torch.no_grad():
    for parameter in zero_mlp.parameters():
        parameter.zero_()
zero_history = fit(zero_mlp, Xtrain, ytrain, Xval, yval, lr=0.1, epochs=200)
print(f"Zero initialization: final training loss {zero_history['train_loss'][-1]:.4f}")
for name, parameter in zero_mlp.named_parameters():
    print(name, "still zero:", bool((parameter == 0).all()))
assert all((zero_mlp[i].weight == 0).all() for i in (0, 2))
assert zero_history["train_loss"][-1] > 1.0

lr=10: largest training loss 218388.78125

Zero initialization: final training loss 1.0511
0.weight still zero: True
0.bias still zero: True
2.weight still zero: True
2.bias still zero: False

Optional: which examples drive the update?

Nothing later in the notebook depends on this section, so you can skip it.

Autodiff is also useful to inspect a model, not only to train it. Fix the parameters and compute one gradient per training example, $g_i=\nabla_\theta\ell_i$. The gradient of the mean training loss is their average, $\bar g=\frac{1}{N}\sum_i g_i$, and a full-batch step moves along $-\bar g$. The individual $g_i$ show which examples pull the parameters, and in which direction.

We use a copy of the lr=0.1 run after 200 updates, whose loss is still far enough from zero to be interesting. We use only training examples, and nothing here changes the model evaluated above.

The simplest approach is a loop over the examples. torch.autograd.grad returns each gradient without touching the .grad fields. We flatten the weight and bias gradients into one vector, which gives a matrix with one row per example and one column per parameter.

analysis_model = copy.deepcopy(lr_models[0.1]).cpu().double().eval()
analysis_model.zero_grad(set_to_none=True)
analysis_X, analysis_y = Xtrain.cpu().double(), ytrain.cpu().clone()
analysis_parameters = tuple(analysis_model.parameters())
analysis_parameter_names = tuple(dict(analysis_model.named_parameters()))

def gradients_in_a_loop(
    Xbatch: Float[torch.Tensor, "batch features"],
    ybatch: Int[torch.Tensor, "batch"],
) -> Float[torch.Tensor, "batch parameters"]:
    rows = []
    for sample, label in zip(Xbatch, ybatch):
        sample_loss = F.cross_entropy(analysis_model(sample[None, :]), label[None])
        derivatives = torch.autograd.grad(sample_loss, analysis_parameters)
        rows.append(torch.cat([value.reshape(-1) for value in derivatives]))
    return torch.stack(rows)

per_example_loop = gradients_in_a_loop(analysis_X, analysis_y)
batch_loss = F.cross_entropy(analysis_model(analysis_X), analysis_y)
batch_derivatives = torch.autograd.grad(batch_loss, analysis_parameters)
mean_gradient = torch.cat([value.reshape(-1) for value in batch_derivatives])
torch.testing.assert_close(per_example_loop.mean(dim=0), mean_gradient)
assert all(parameter.grad is None for parameter in analysis_parameters)
print("Gradient matrix (examples, parameters):", tuple(per_example_loop.shape))

Gradient matrix (examples, parameters): (205, 15)

The same computation with function transforms

torch.func lets us write the gradient for one example and apply it to the whole batch:

• functional_call(model, params, inputs) runs the model with an explicit dictionary of parameters;
• grad(f) returns a function that computes the gradient of f with respect to its first argument;
• vmap(f, in_dims=(None, 0, 0)) applies f to each example. None means that the argument (here, the parameters) is shared; 0 means that we map over the first axis of the argument.

The result has an extra leading axis for the examples. After concatenating, it must match the loop. This works directly because our model has no dropout, batch normalization or other layers that depend on the whole batch; for such models, a one-example loss may not describe what the model computes on a batch. We do not compare timings for this small model. See PyTorch's per-sample-gradient tutorial.

from torch.func import functional_call, grad, vmap

analysis_state = {name: value.detach() for name, value in analysis_model.named_parameters()}

def single_example_loss(
    parameters: dict[str, torch.Tensor],
    sample: Float[torch.Tensor, "features"],
    label: Int[torch.Tensor, ""],
) -> Float[torch.Tensor, ""]:
    sample_logits = functional_call(analysis_model, parameters, (sample.unsqueeze(0),))
    return F.cross_entropy(sample_logits, label.unsqueeze(0))

one_gradient = grad(single_example_loss)
all_gradients = vmap(one_gradient, in_dims=(None, 0, 0))
gradient_tree = all_gradients(analysis_state, analysis_X, analysis_y)
per_example_transformed = torch.cat([
    gradient_tree[name].reshape(len(analysis_y), -1) for name in analysis_parameter_names
], dim=1)
torch.testing.assert_close(per_example_transformed, per_example_loop)
print("The loop and grad + vmap agree.")

The loop and grad + vmap agree.

Magnitude and direction

Two numbers summarize each $g_i$. The norm $\lVert g_i\rVert$ measures how strongly this example pulls on the parameters. The cosine $g_i^\top\bar g/(\lVert g_i\rVert\lVert\bar g\rVert)$ measures whether it agrees with the average. If the cosine is positive, a small enough step along $-\bar g$ lowers this example's loss; if it is negative, the step raises it. This holds only locally, for small steps from the current parameters, and the cosine is undefined when either vector is zero.

We highlight an easy example (correctly classified, with low loss) and a difficult one (the highest loss). Then we take the easy example and change only its label, keeping the model fixed. Before running, predict where this relabeled example will land in the two plots.

with torch.no_grad():
    clean_logits = analysis_model(analysis_X)
    clean_losses = F.cross_entropy(clean_logits, analysis_y, reduction="none")
correct_indices = torch.where(clean_logits.argmax(dim=1) == analysis_y)[0]
ordered_correct = correct_indices[clean_losses[correct_indices].argsort()]
easy_index = int(ordered_correct[len(ordered_correct) // 10])
difficult_index = int(clean_losses.argmax())
changed_label = (analysis_y[easy_index:easy_index + 1] + 1) % len(classes)
changed_gradient = gradients_in_a_loop(analysis_X[easy_index:easy_index + 1], changed_label)[0]
with torch.no_grad():
    changed_loss = F.cross_entropy(clean_logits[easy_index:easy_index + 1], changed_label)

def gradient_alignment(
    vectors: Float[torch.Tensor, "batch parameters"],
    direction: Float[torch.Tensor, "parameters"],
) -> Float[torch.Tensor, "batch"]:
    denominator = vectors.norm(dim=1) * direction.norm()
    result = torch.full_like(denominator, float("nan"))
    valid = denominator > 1e-12
    result[valid] = (vectors[valid] @ direction) / denominator[valid]
    return result

clean_norms = per_example_loop.norm(dim=1)
clean_alignment = gradient_alignment(per_example_loop, mean_gradient)
comparison_gradients = torch.stack([
    per_example_loop[easy_index], per_example_loop[difficult_index], changed_gradient
])
comparison = pd.DataFrame({
    "case": ["easy", "difficult", "easy, changed label"],
    "training row": [easy_index, difficult_index, easy_index],
    "label": [classes[int(analysis_y[easy_index])], classes[int(analysis_y[difficult_index])],
              classes[int(changed_label[0])]],
    "loss": [clean_losses[easy_index].item(), clean_losses[difficult_index].item(), changed_loss.item()],
    "gradient norm": comparison_gradients.norm(dim=1).numpy(),
    "cosine with clean mean": gradient_alignment(comparison_gradients, mean_gradient).numpy(),
})
print(comparison.to_string(index=False, float_format=lambda value: f"{value:.4f}"))

fig, axes = plt.subplots(1, 2, figsize=(11, 4))
axes[0].scatter(clean_losses.numpy(), clean_norms.numpy(), s=18, alpha=0.35, label="Original labels")
axes[1].scatter(clean_norms.numpy(), clean_alignment.numpy(), s=18, alpha=0.35)
for row, marker, color in zip(comparison.to_dict("records"), ["o", "s", "X"], ["C2", "C1", "C3"]):
    axes[0].scatter(row["loss"], row["gradient norm"], marker=marker, color=color, s=85, label=row["case"])
    axes[1].scatter(row["gradient norm"], row["cosine with clean mean"], marker=marker, color=color, s=85)
axes[0].set(xlabel="Per-example cross-entropy", ylabel="Gradient norm")
axes[1].set(xlabel="Gradient norm", ylabel="Cosine with clean mean gradient", ylim=(-1.05, 1.05))
axes[1].axhline(0, color="gray", lw=1)
axes[0].legend(fontsize=8)
fig.tight_layout()
plt.show()

               case  training row     label   loss  gradient norm  cosine with clean mean
               easy            84    Gentoo 0.0084         0.0333                  0.2910
          difficult            22 Chinstrap 1.4234         1.9860                  0.0718
easy, changed label            84    Adelie 7.3561         4.0911                 -0.3475

Which examples dominate the mean gradient: those with a large norm, those aligned with the mean, or both? Can large gradients cancel each other out?

A large gradient norm does not mean that an example is mislabeled. What else could produce one? Think about ambiguous measurements, the limits of a linear model and the current parameter values. Repeat the analysis with another run from the learning-rate comparison: do the same examples stand out?

Instructor discussion

The loop and the transforms compute the same derivatives, and the check against the mean links the per-example view back to the full-batch update. torch.func.grad differentiates with respect to the explicit parameter dictionary, even though its values are detached, and it does not fill the module's .grad fields.

Changing only the label changes both the size and the direction of the gradient. Still, a large norm is not a test of label quality: a correctly labeled but hard example can also disagree strongly with the current model. The cosine is always relative to the mean gradient. Near a minimum the mean gets small and its direction becomes unreliable. Measuring how much one example influences the final trained model would need retraining, which is a different experiment.

Optional exercise (after the MLP lecture): store an image in an MLP

Section 9 swapped the linear model for an MLP on the penguins. Here we take the same tools to a different problem: a regression with a different loss and input. We fit a function that maps a pixel position (x, y) to its color (r, g, b), so that the network itself stores the image. The next cell generates a small target image, so no download is needed.

1. Build an MLP with two ReLU hidden layers and three sigmoid outputs. Train it with the mean squared error, using only the pixels in image_train_mask (one pixel in four). Coordinates are in [-1, 1] and colors in [0, 1].
2. Compare two inputs: the raw coordinates, and Fourier features, the concatenation of sin(2*pi*coordinates @ frequencies.T) and cos(...). Here frequencies is a fixed matrix of 16 two-dimensional vectors, drawn once from a normal distribution with standard deviation 1.5 and never trained.
3. For each input, try hidden widths 16 and 32, with seed 7, Adam with learning rate 0.01 and 400 full-batch updates. Record the error on the training pixels and on the held-out pixels, and save the reconstructions at updates 0, 100 and 400. Fix these settings before you look at held-out results.
4. Evaluate the network on a finer 64 by 64 grid. Does the larger image contain detail that the network could not have seen?

The held-out pixels test interpolation inside this one image, not generalization to new images. If you tune your choices on them, keep another set of pixels for the final evaluation.

If you want to go further: plot the held-out error against the number of trainable parameters (at the same width, the two inputs give different counts), and compare the float32 size of the network, including the frequency matrix, with the original image at one byte per channel. What does this comparison leave out? You can also query coordinates outside [-1, 1], where the network is extrapolating.

The instructor version contains a reference solution. Background: Tancik et al., Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains.

image_side = 32

def coordinate_grid(side: int) -> Float[torch.Tensor, "pixels 2"]:
    axis = (torch.arange(side, dtype=torch.float32) + 0.5) * (2. / side) - 1.
    yy, xx = torch.meshgrid(axis, axis, indexing="ij")
    return torch.stack([xx, yy], dim=-1).reshape(-1, 2)

image_coordinates = coordinate_grid(image_side)
image_x, image_y = image_coordinates.unbind(dim=1)
# A small RGB test image with smooth variation, stripes and a circular edge.
image_rgb = torch.stack([
    0.5 + 0.4 * torch.sin(3 * torch.pi * image_x),
    0.5 + 0.4 * torch.cos(2 * torch.pi * image_y),
    torch.where(image_x.square() + image_y.square() < 0.45, 0.85, 0.15),
], dim=1)
image_bytes = (255 * image_rgb).round().to(torch.uint8)
image_rgb = image_bytes.float() / 255.
pixel_rows, pixel_cols = torch.meshgrid(torch.arange(image_side), torch.arange(image_side), indexing="ij")
image_train_mask = ((pixel_rows % 2 == 0) & (pixel_cols % 2 == 0)).flatten()
image_heldout_mask = ~image_train_mask
assert image_train_mask.sum().item() == 256 and image_heldout_mask.sum().item() == 768

fig, axes = plt.subplots(1, 2, figsize=(6, 3))
axes[0].imshow(image_rgb.reshape(image_side, image_side, 3))
axes[0].set_title("Target RGB image")
visible = torch.ones_like(image_rgb)
visible[image_train_mask] = image_rgb[image_train_mask]
axes[1].imshow(visible.reshape(image_side, image_side, 3))
axes[1].set_title("Observed pixels (others white)")
for ax in axes:
    ax.axis("off")
fig.tight_layout()
plt.show()

# Define the coordinate MLP, train on the observed pixels, and compare encodings.
# Keep this optional exercise independent of the penguin classifier above.

Instructor reference experiment

The four configurations use the prescribed settings and run on CPU. The frequency matrix is a registered buffer: the optimizer does not update it, but it is saved with the model. The network never sees the formula that generated the image, only pixel values.

This is a small teaching run with one seed, not a tuned comparison of encodings or widths.

class CoordinateMLP(torch.nn.Module):
    def __init__(self, width: int, fourier: bool):
        super().__init__()
        frequency_generator = torch.Generator().manual_seed(7)
        frequencies = 1.5 * torch.randn(16, 2, generator=frequency_generator) if fourier else torch.empty(0, 2)
        self.register_buffer("frequencies", frequencies)
        input_features = 2 * len(frequencies) if fourier else 2
        self.layers = torch.nn.Sequential(
            torch.nn.Linear(input_features, width), torch.nn.ReLU(),
            torch.nn.Linear(width, width), torch.nn.ReLU(),
            torch.nn.Linear(width, 3), torch.nn.Sigmoid(),
        )

    def forward(self, coordinates: Float[torch.Tensor, "pixels 2"]) -> Float[torch.Tensor, "pixels 3"]:
        if len(self.frequencies):
            phase = 2 * torch.pi * coordinates @ self.frequencies.T
            coordinates = torch.cat([phase.sin(), phase.cos()], dim=1)
        return self.layers(coordinates)


def fit_coordinate_mlp(width: int, fourier: bool):
    with torch.random.fork_rng(devices=[]):
        torch.manual_seed(7)
        coordinate_model = CoordinateMLP(width, fourier)
    optimizer = torch.optim.Adam(coordinate_model.parameters(), lr=0.01)
    trace = {"update": [], "train_mse": [], "heldout_mse": []}
    snapshots = {}
    for update in range(401):
        if update % 20 == 0:
            with torch.no_grad():
                reconstructed = coordinate_model(image_coordinates)
                trace["update"].append(update)
                trace["train_mse"].append(F.mse_loss(reconstructed[image_train_mask], image_rgb[image_train_mask]).item())
                trace["heldout_mse"].append(F.mse_loss(reconstructed[image_heldout_mask], image_rgb[image_heldout_mask]).item())
                if update in (0, 100, 400):
                    snapshots[update] = reconstructed.reshape(image_side, image_side, 3).clone()
        if update == 400:
            break
        optimizer.zero_grad(set_to_none=True)
        prediction = coordinate_model(image_coordinates[image_train_mask])
        image_loss = F.mse_loss(prediction, image_rgb[image_train_mask])
        image_loss.backward()
        optimizer.step()
    return coordinate_model, trace, snapshots

image_runs = {}
for encoding in ("raw", "Fourier"):
    for width in (16, 32):
        image_runs[(encoding, width)] = fit_coordinate_mlp(width, encoding == "Fourier")

image_summary_rows = []
for (encoding, width), (coordinate_model, trace, snapshots) in image_runs.items():
    parameter_count = sum(p.numel() for p in coordinate_model.parameters())
    payload_bytes = sum(value.numel() * value.element_size() for value in coordinate_model.state_dict().values())
    image_summary_rows.append({
        "encoding": encoding, "width": width, "parameters": parameter_count,
        "float32 payload bytes": payload_bytes,
        "train MSE": trace["train_mse"][-1], "held-out MSE": trace["heldout_mse"][-1],
    })
image_summary = pd.DataFrame(image_summary_rows)
print(image_summary.to_string(index=False, float_format=lambda value: f"{value:.6f}"))
print("Original uint8 RGB payload bytes:", image_bytes.numel() * image_bytes.element_size())

fig, axes = plt.subplots(2, 3, figsize=(8, 5.5))
for row, encoding in enumerate(("raw", "Fourier")):
    for col, update in enumerate((0, 100, 400)):
        axes[row, col].imshow(image_runs[(encoding, 32)][2][update])
        axes[row, col].set_title(f"{encoding}, width 32, step {update}", fontsize=10)
        axes[row, col].axis("off")
fig.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 3, figsize=(13, 3.5))
for (encoding, width), (_, trace, _) in image_runs.items():
    label = f"{encoding}, {width}"
    style = "-" if encoding == "raw" else "--"
    axes[0].plot(trace["update"], trace["train_mse"], style, label=label)
    axes[1].plot(trace["update"], trace["heldout_mse"], style, label=label)
for row in image_summary.to_dict("records"):
    axes[2].scatter(row["parameters"], row["held-out MSE"], label=f"{row['encoding']}, {row['width']}")
axes[0].set(xlabel="Update", ylabel="Training pixel MSE", yscale="log")
axes[1].set(xlabel="Update", ylabel="Held-out pixel MSE", yscale="log")
axes[2].set(xlabel="Trainable parameters", ylabel="Held-out pixel MSE", yscale="log")
axes[0].legend(fontsize=8)
axes[2].legend(fontsize=8)
fig.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 2, figsize=(6, 3))
with torch.no_grad():
    for ax, encoding in zip(axes, ("raw", "Fourier")):
        dense_rgb = image_runs[(encoding, 32)][0](coordinate_grid(64)).reshape(64, 64, 3)
        ax.imshow(dense_rgb)
        ax.set_title(f"{encoding}: 64 by 64 query")
        ax.axis("off")
fig.tight_layout()
plt.show()

encoding  width  parameters  float32 payload bytes  train MSE  held-out MSE
     raw     16         371                   1484   0.016235      0.017249
     raw     32        1251                   5004   0.006098      0.008555
 Fourier     16         851                   3532   0.000829      0.029108
 Fourier     32        2211                   8972   0.000064      0.017427
Original uint8 RGB payload bytes: 3072

Reading the image experiment

The network sees only a quarter of the pixels; the others measure how well it interpolates. In the recorded CPU run, Fourier features fit the training pixels much more closely, while the raw-coordinate model with width 32 has the lowest held-out error. Fitting the observed pixels and filling in the missing ones are different things, and the training error alone cannot tell them apart. Compare the images and the held-out curves before giving Fourier features the credit for any improvement.

The storage column counts every parameter and buffer, including the frequency matrix. It leaves out the code needed to decode the image, the architecture description and any file format overhead. It compares uncompressed sizes, not PNG or JPEG files, so even a network smaller than the raw image would not be a compression result without a codec comparison at equal quality.

The 64 by 64 render evaluates the same function at more points; it adds no information, and outside the training range the function is even less constrained. Other seeds, images or Fourier scales may change the ranking. These are good follow-up questions, not conclusions of this four-run example.

References

Data: Horst, Hill and Gorman (2020), palmerpenguins, CC0. Original data collected by Kristen Gorman and Palmer Station LTER; see the dataset page for the ecological study and data citations.

• CrossEntropyLoss
• nn.Module
• Autograd mechanics
• Per-sample gradients with torch.func
• Fourier feature networks