The Network: What Is It Without nn.Module?
Here is a two-layer neural network, written with nothing but the machinery of the last three chapters. Four trainable tensors, a forward pass, a loss, backward(), an update. It runs without warnings, and the loss falls by a third.
import math
import torch
torch.manual_seed(0)
n = 2000
X = torch.randn(n, 2)
y = ((X[:, 0] > 0) ^ (X[:, 1] > 0)).long()
perm = torch.randperm(n)
n_train = int(0.8 * n)
X_train, y_train = X[perm[:n_train]], y[perm[:n_train]]
X_val, y_val = X[perm[n_train:]], y[perm[n_train:]]
D_in, H, C = 2, 16, 2
torch.manual_seed(1)
W1 = (torch.randn(D_in, H) / math.sqrt(D_in)).requires_grad_()
b1 = torch.zeros(H, requires_grad=True)
W2 = (torch.randn(H, C) / math.sqrt(H)).requires_grad_()
b2 = torch.zeros(C, requires_grad=True)
parameters = [W1, b1, W2, b2]
def forward(x):
z1 = x @ W1 + b1
return z1 @ W2 + b2
def cross_entropy(logits, targets):
log_probs = torch.log_softmax(logits, dim=1)
return -log_probs[torch.arange(targets.shape[0]), targets].mean()
def accuracy(logits, targets):
return (logits.argmax(dim=1) == targets).float().mean().item()
learning_rate = 0.5
# First verify that one training step is mechanically healthy.
names = ["W1", "b1", "W2", "b2"]
initial_values = [p.detach().clone() for p in parameters]
loss = cross_entropy(forward(X_train), y_train)
before = [p.detach().clone() for p in parameters]
loss.backward()
with torch.no_grad():
for p in parameters:
p -= learning_rate * p.grad
for name, p, old in zip(names, parameters, before):
print(f"{name:3s} leaf={p.is_leaf} grad_none={p.grad is None} "
f"grad_norm={p.grad.norm().item():8.4f} "
f"finite={torch.isfinite(p.grad).all().item()} "
f"moved={(p.detach() - old).abs().max().item():.4f}")
W1 leaf=True grad_none=False grad_norm= 0.4516 finite=True moved=0.1048
b1 leaf=True grad_none=False grad_norm= 0.0192 finite=True moved=0.0053
W2 leaf=True grad_none=False grad_norm= 1.1058 finite=True moved=0.1918
b2 leaf=True grad_none=False grad_norm= 0.0175 finite=True moved=0.0062
Every parameter is a leaf. Every parameter requires gradients. Every parameter receives a finite gradient, and every parameter changes by a measurable amount when the update runs. The basic training machinery is working.
Now restore the exact starting state so that the diagnostic step does not alter the experiment:
with torch.no_grad():
for p, initial in zip(parameters, initial_values):
p.copy_(initial)
for p in parameters:
p.grad = None
Then run the full training experiment:
for step in range(2001):
logits = forward(X_train)
loss = cross_entropy(logits, y_train)
if step % 500 == 0:
with torch.no_grad():
print(f"step={step:4d} loss={loss.item():.4f} "
f"train_acc={accuracy(logits, y_train):.3f} "
f"val_acc={accuracy(forward(X_val), y_val):.3f}")
for p in parameters:
if p.grad is not None:
p.grad.zero_()
loss.backward()
with torch.no_grad():
for p in parameters:
p -= learning_rate * p.grad
step= 0 loss=1.0665 train_acc=0.506 val_acc=0.502
step= 500 loss=0.6928 train_acc=0.596 val_acc=0.592
step=1000 loss=0.6928 train_acc=0.596 val_acc=0.592
step=1500 loss=0.6928 train_acc=0.596 val_acc=0.592
step=2000 loss=0.6928 train_acc=0.596 val_acc=0.592
The loss dropped from 1.0665 to 0.6928 and then stopped, permanently, at a value suspiciously close to log 2. Accuracy settled just under sixty percent on a two-class problem.
The one-step diagnostic already established that the training machinery works: gradients exist, they are finite, and the parameters update. Yet the full run still fails to solve the problem.
The missing question is therefore not merely whether the implementation executes correctly. It is what function this particular arrangement of tensors is capable of representing.
The health checks we have used so far do not explain the failure. The tensors have legal shapes, the autograd paths exist, the gradients are finite, and the parameters update. The missing question is not whether the implementation executes correctly, but what function this architecture is capable of representing. It is in what this particular arrangement of tensors is capable of computing, and that is a question we have not yet learned to ask.
What is a neural network, once you remove
nn.Module,nn.Linearand the optimizer? Which piece of it is actually doing the work, and how do you tell a broken implementation apart from an architecture that cannot represent the task?
By the end of this chapter you will have built that network from raw tensors, understood what each part contributes, verified your hand-written loss against PyTorch’s implementation on both value and gradient, and collected the specific evidence that distinguishes the failure above from the four other things it superficially resembles.
The dataset is a measuring instrument
Look at what the data actually is.
y = ((X[:, 0] > 0) ^ (X[:, 1] > 0)).long()
The label is 1 when the two coordinates have opposite signs and 0 when they agree. That is XOR on the signs of the inputs: class 1 lives in the second and fourth quadrants, class 0 in the first and third.
This dataset was not chosen because it is realistic. It was chosen because it has a property that makes it useful as an instrument: no single straight line separates the two classes. Whichever half-plane you pick, it contains points of both classes, because the two class regions meet at the origin from opposite sides. That is a fact about the geometry of the problem, and it holds no matter how the model is implemented, how the optimizer is tuned, or how long you train.
Which means we can measure something normally invisible. Instead of arguing vaguely about whether a linear model could do better, we can search a dense grid of linear decision boundaries and obtain an empirical reference for what this model class can achieve on this finite dataset.
def best_linear_accuracy(Xs, ys, n_angles=720, n_offsets=601):
"""Search a dense grid of half-plane classifiers and return the best accuracy found."""
angles = torch.linspace(0, math.pi, n_angles + 1)[:-1]
offsets = torch.linspace(-3, 3, n_offsets)
best = 0.0
for a in angles:
w = torch.tensor([a.cos(), a.sin()])
scores = Xs @ w # (N,)
preds = scores.unsqueeze(1) > offsets.unsqueeze(0) # (N, n_offsets)
acc = (preds.long() == ys.unsqueeze(1)).float().mean(0)
best = max(best, acc.max().item(), (1 - acc).max().item())
return best
print("best grid-search linear train accuracy:", round(best_linear_accuracy(X_train, y_train), 4))
best grid-search linear train accuracy: 0.6431
The dense search finds linear classifiers around 64% accuracy on this particular training sample. Our trained network reaches 59.6%, so there is still an optimization gap between the model we obtained and the best boundary found by the grid search.
That does not weaken the architectural result. Even if we closed that gap completely, an affine classifier still could not represent the XOR rule. Better optimization can move us within the affine model class; it cannot move us outside it.
That number is worth taking seriously as a piece of evidence, because it belongs to a category the earlier chapters never produced:
An independent reference for a model class helps separate optimization from representation. If a dense search finds affine classifiers around 64% on this dataset and your trained affine model reaches 59%, there may still be an optimization gap to close. But no amount of optimization can make an affine model represent the XOR rule exactly.
But the network has two layers. It has sixteen hidden units. Why is it behaving like a straight line?
Two affine layers are one affine layer
Write the forward pass out symbolically, with X of shape (B, 2):
z1 = X @ W1 + b1
logits = z1 @ W2 + b2
Substitute the first line into the second and expand:
logits = (X @ W1 + b1) @ W2 + b2
= X @ W1 @ W2 + b1 @ W2 + b2
= X @ (W1 @ W2) + (b1 @ W2 + b2)
Matrix multiplication is associative, so X @ W1 @ W2 regroups freely, and W1 @ W2 is a single matrix of shape (2, 2). The bias terms collapse into a single vector of shape (2,). The entire two-layer network is identically a one-layer network:
logits = X @ W_eq + b_eq W_eq = W1 @ W2 b_eq = b1 @ W2 + b2
This is a claim about exact equality, not an approximation, so check it as one.
with torch.no_grad():
W_eq = W1 @ W2
b_eq = b1 @ W2 + b2
two_layer = forward(X_train)
one_layer = X_train @ W_eq + b_eq
print("W_eq shape:", tuple(W_eq.shape), " b_eq shape:", tuple(b_eq.shape))
print("max abs difference:", (two_layer - one_layer).abs().max().item())
print("parameters in the two-layer form:", sum(p.numel() for p in parameters))
print("parameters in the collapsed form:", W_eq.numel() + b_eq.numel())
Run this on the network after the two thousand training steps above:
W_eq shape: (2, 2) b_eq shape: (2,)
max abs difference: 7.152557373046875e-07
parameters in the two-layer form: 82
parameters in the collapsed form: 6
The difference is at the float32 noise floor. Eighty-two trainable numbers were used to express logits that six numbers can describe exactly. The extra seventy-six numbers do not enlarge the function class; they re-parameterize the same affine map.
This proves a representational limitation, not that optimization found the best affine solution. The 59.6% result belongs to the particular solution gradient descent reached; the dense linear search already showed that somewhat better affine classifiers exist on this sample.
The same check on the untrained network gives 1.43e-06, so this is not something training produced. It was true from the moment the parameters were created.
A test that does not require reading the source
The collapse above needed access to W1 and W2. Often you have a model object, a stack of generated code, or a forward you would rather not read line by line. There is a way to ask the same question from outside.
An affine function satisfies an additivity identity. If f(x) = x @ W + b, then for any two inputs:
f(x1 + x2) - f(0) = (f(x1) - f(0)) + (f(x2) - f(0))
Subtracting f(0) removes the bias, leaving the linear part, and linear maps distribute over addition. Any function that fails this identity is not affine. Passing the identity on a finite collection of probes is weaker evidence: it is consistent with the function being affine, but it does not prove that the function is affine everywhere. A nonlinear function can behave affinely over the region you happened to probe.
def affine_residual(f, x1, x2):
"""Measure affine-identity violation on these probes. A non-zero residual falsifies affinity."""
with torch.no_grad():
zero = torch.zeros_like(x1)
lhs = f(x1 + x2) - f(zero)
rhs = (f(x1) - f(zero)) + (f(x2) - f(zero))
return (lhs - rhs).abs().max().item()
a = torch.randn(64, 2) # random probe inputs; any two batches will do
b = torch.randn(64, 2)
print("no activation:", affine_residual(lambda x: (x @ W1 + b1) @ W2 + b2, a, b))
print("with ReLU: ", affine_residual(lambda x: torch.relu(x @ W1 + b1) @ W2 + b2, a, b))
no activation: 9.5367431640625e-07
with ReLU: 2.2889842987060547
The second number depends on which probes you happen to draw and on the current parameter values, so do not read anything into its exact size. The comparison that matters is 1e-6 against order one: one of these functions is affine to float precision and the other is not.
Two lines of arithmetic can quickly prove that a callable is not affine. If the residual stays near zero across many well-chosen probes, that is useful evidence of affine behavior, but the algebraic collapse above remains the proof when the implementation is available. This is worth keeping. When someone hands you a model with a promising-looking architecture diagram and it will not beat a linear baseline, a large residual can rule out the explanation that the callable is affine. A near-zero residual keeps that explanation plausible and tells you where to investigate next; it does not prove affinity by itself.
The repair, and the evidence for it
Change one thing in the original program. Wrap the first layer’s output in torch.relu before it reaches the second layer:
def forward(x):
z1 = x @ W1 + b1
h = torch.relu(z1)
return h @ W2 + b2
Same data, same seeds, same initial parameter values, same learning rate, same number of steps.
step= 0 loss=0.8268 train_acc=0.281 val_acc=0.278
step= 500 loss=0.0589 train_acc=0.994 val_acc=0.995
step=1000 loss=0.0408 train_acc=0.996 val_acc=0.995
step=1500 loss=0.0327 train_acc=0.997 val_acc=0.995
step=2000 loss=0.0279 train_acc=0.998 val_acc=0.995
99.5% on held-out data. The affine predecessor could not represent the XOR rule at all, while the nonlinear model can. The dense linear search gave us an empirical reference around 64% on this training sample; the important difference is not that 64% is an exact ceiling, but that the affine function class cannot express the required boundary.
Be precise about what that demonstrates, because this is exactly the kind of result that attracts sloppy summaries. It does not show that ReLU is the best activation, that nonlinearities make optimization easier, or that deeper is better. What it shows is narrower and more useful:
Stacking affine transformations without a nonlinearity does not enlarge the function class: the composition collapses to another affine transformation. A nonlinear activation prevents that algebraic collapse, allowing additional layers to represent functions that a single affine map cannot.
torch.relu(z) is elementwise max(z, 0). It is not affine, it does not commute with matrix multiplication, and so relu(X @ W1 + b1) @ W2 + b2 cannot be regrouped into anything simpler. Each hidden ReLU unit introduces a potential hinge along the boundary where its pre-activation crosses zero. With sixteen hidden units, the model has up to sixteen such boundaries available, although some may coincide, lie outside the relevant data region, or contribute nothing after the second layer. That combination is piecewise linear, and a piecewise linear function with enough pieces can trace the boundary between opposite quadrants.
There is one more thing worth noticing about the run above, and it will matter for the rest of the chapter. The failing model’s loss also went down. It fell by a third, monotonically, exactly as a working model’s loss does. The loss curve did not distinguish the two situations. Held-out accuracy showed that the model was performing poorly, while the dense linear search helped tell us that the poor result was consistent with the limitations of an affine model.
The parameters are the model
We have repaired the opening program. We have not yet built it. Now do that deliberately, from the bottom, connecting each line back to the chapter that explains it.
Start with the only thing in the system the training loop is allowed to modify.
D_in, H, C = 2, 16, 2
torch.manual_seed(1)
W1 = (torch.randn(D_in, H) / math.sqrt(D_in)).requires_grad_()
b1 = torch.zeros(H, requires_grad=True)
W2 = (torch.randn(H, C) / math.sqrt(H)).requires_grad_()
b2 = torch.zeros(C, requires_grad=True)
parameters = [W1, b1, W2, b2]
for name, p in zip(["W1", "b1", "W2", "b2"], parameters):
print(f"{name:3s} shape={str(tuple(p.shape)):9s} numel={p.numel():3d} "
f"requires_grad={p.requires_grad} is_leaf={p.is_leaf} grad_fn={p.grad_fn}")
print("total trainable numbers:", sum(p.numel() for p in parameters))
W1 shape=(2, 16) numel= 32 requires_grad=True is_leaf=True grad_fn=None
b1 shape=(16,) numel= 16 requires_grad=True is_leaf=True grad_fn=None
W2 shape=(16, 2) numel= 32 requires_grad=True is_leaf=True grad_fn=None
b2 shape=(2,) numel= 2 requires_grad=True is_leaf=True grad_fn=None
total trainable numbers: 82
Those four tensors are the model. There is no other object. When the chapter title asks what a neural network is without nn.Module, this is the answer: eighty-two numbers arranged into four tensors, plus a function that says how to combine them with an input.
Chapter 2’s discipline applies directly. Each axis means something, and nothing in PyTorch enforces it:
W1 (D_in, H) axis 0 indexes input features, axis 1 indexes hidden units
b1 (H,) one offset per hidden unit
W2 (H, C) axis 0 indexes hidden units, axis 1 indexes output classes
b2 (C,) one offset per class
Read W1[i, j] as how strongly input feature i drives hidden unit j. Read a column W1[:, j] as the entire incoming weight vector of hidden unit j. Those two readings are the ones you will need when we ask why identically-initialized units stay identical.
Chapter 3’s discipline applies too, and it is the reason for the slightly awkward spelling of the two weight lines.
The initialization that quietly stops being a parameter
These two lines look like the same thing said two ways:
W1 = torch.randn(2, 16, requires_grad=True) / math.sqrt(2) # A
W1 = (torch.randn(2, 16) / math.sqrt(2)).requires_grad_() # B
They are not. In A, requires_grad=True is set on the output of randn, and then that tensor is divided by a constant. Division is a recorded operation, so the name W1 ends up bound to the result of a recorded operation. In B, the division happens first on an ordinary tensor, and requires_grad_() flags the finished result.
Do not take that on trust. Chapter 3 gave you three fields that settle it:
A = torch.randn(2, 16, requires_grad=True) / math.sqrt(2)
B = (torch.randn(2, 16) / math.sqrt(2)).requires_grad_()
for name, t in (("A", A), ("B", B)):
print(f"{name}: requires_grad={t.requires_grad} is_leaf={t.is_leaf} "
f"grad_fn={type(t.grad_fn).__name__ if t.grad_fn is not None else None}")
A: requires_grad=True is_leaf=False grad_fn=DivBackward0
B: requires_grad=True is_leaf=True grad_fn=None
A requires gradients, which is the property people check, and is not a leaf, which is the property that matters for manually managed parameters. Predict the consequences before reading on. A is not a leaf, so backward() will not populate A.grad — that is documented behavior, not a bug. And A carries a live grad_fn, so the division node is part of the graph on every forward pass.
Build the network with W1 created the wrong way and run it:
torch.manual_seed(1)
W1 = torch.randn(D_in, H, requires_grad=True) / math.sqrt(D_in) # not a leaf
b1 = torch.zeros(H, requires_grad=True)
W2 = (torch.randn(H, C) / math.sqrt(H)).requires_grad_()
b2 = torch.zeros(C, requires_grad=True)
step= 0 loss=0.8268 train_acc=0.281 val_acc=0.278
RuntimeError: Trying to backward through the graph a second time (or directly
access saved tensors after they have already been freed). Saved intermediate
values of the graph are freed when you call .backward() or autograd.grad().
Specify retain_graph=True if you need to backward through the graph a second
time or ...
Step 0 completes. Step 1 raises, and the exception names a line that has nothing to do with the mistake. This is Chapter 3’s graph-lifetime mechanism arriving in a new disguise: W1 is a non-leaf whose DivBackward0 node was created once, outside the loop, and every iteration builds a fresh graph that terminates in that same reused node. The first backward() frees it. The second finds it already gone.
The error message contains a suggestion, and the suggestion works, in the sense that the exception disappears:
loss.backward(retain_graph=True)
step= 0 loss=0.8268 train_acc=0.281 val_acc=0.278
step= 100 loss=0.2695 train_acc=0.936 val_acc=0.950
step= 500 loss=0.3011 train_acc=0.836 val_acc=0.837
step=1000 loss=0.2988 train_acc=0.837 val_acc=0.850
step=2000 loss=0.2879 train_acc=0.842 val_acc=0.858
print("W1.grad is None:", W1.grad is None)
print("W1 unchanged after 2000 steps:", torch.equal(W1_before, W1.detach()))
W1.grad is None: True
W1 unchanged after 2000 steps: True
The program now runs to completion and reaches 86% validation accuracy, which is far enough above the linear ceiling to look like success. W1 was never trained. It holds its random initial values for all two thousand steps, and the guarded update loop skipped it silently every time because p.grad is not None was false. What actually trained was b1, W2 and b2 on top of sixteen frozen random features, which is a real model, just not the one anyone intended.
PyTorch does emit a UserWarning the first time you read .grad on a non-leaf, and that warning is the honest signal here. It is also easy to lose in training output, and it fires at the inspection site rather than at the creation site.
Three points are worth extracting, and they generalize well beyond this line of code.
The first is the diagnostic move: the exception was raised at loss.backward() on step 1, and the mistake was on the parameter creation line before the loop. Chapter 2 taught us to look for the first wrong tensor rather than the first illegal operation; Chapter 3 taught us to look for the place where the path breaks rather than the place where .grad is missing. Here both apply at once.
The second is about the suggested fix. retain_graph=True is a real feature with real uses, and it removed the symptom completely. It also converted a loud failure into a silent one. When an error message proposes a remedy, the remedy addresses the exception, not necessarily your intent.
The third is about what to check. requires_grad was True throughout. Only is_leaf distinguished the two cases, and only a snapshot comparison proved that a parameter had stopped moving.
For manually managed trainable state,
requires_grad=Trueandis_leaf=Trueare two separate requirements. Create the tensor first, then flag it.
nn.Parameter exists partly to make this class of mistake harder to commit, and Chapter 5 is where that belongs. For now we are managing the leaves ourselves, which is precisely why we have to know what a leaf is.
What a layer is
With the parameters in hand, the forward pass is three lines of tensor arithmetic.
def forward(x):
z1 = x @ W1 + b1 # (B, D_in) @ (D_in, H) -> (B, H)
h = torch.relu(z1) # (B, H)
return h @ W2 + b2 # (B, H) @ (H, C) -> (B, C)
There is nothing mysterious inside the fully connected linear transformation: it is the affine operation x @ W + b. PyTorch later packages that pattern as nn.Linear.
Not every neural-network layer is affine — ReLU, normalization, convolution, attention and many other operations are also called layers — so the useful point here is narrower: this particular linear layer is just tensor arithmetic plus trainable state. Trace the shapes on a small batch and confirm each step against what you predicted:
xb = X_train[:8]
z1 = xb @ W1
h = torch.relu(z1 + b1)
print("x ", tuple(xb.shape))
print("x @ W1 ", tuple(z1.shape))
print("+ b1 ", tuple((z1 + b1).shape), " b1 is", tuple(b1.shape))
print("relu ", tuple(h.shape))
print("h @ W2 ", tuple((h @ W2).shape))
print("+ b2 ", tuple((h @ W2 + b2).shape))
x (8, 2)
x @ W1 (8, 16)
+ b1 (8, 16) b1 is (16,)
relu (8, 16)
h @ W2 (8, 2)
+ b2 (8, 2)
Two details deserve naming rather than being absorbed silently.
The bias broadcast. z1 is (8, 16) and b1 is (16,). Chapter 2’s rule aligns shapes from the right and inserts leading dimensions of size one, so b1 is treated as (1, 16) and expanded down the batch. Every example receives the same sixteen offsets, one per hidden unit. That is the intended semantics: a bias belongs to a unit, not to an example.
The failure worth guarding against is not a shape error, because a shape error would announce itself. It is a bias whose shape is legal and means the wrong thing. A (B, 1) bias, which is what you get if you build offsets per example rather than per unit, broadcasts just as happily:
print(tuple((torch.zeros(8, 16) + torch.zeros(8, 1)).shape))
(8, 16)
Same output shape, completely different model. Two assertions cost nothing and pin the meaning down:
assert b1.shape == (H,)
assert z1.shape == (xb.shape[0], H)
The batch axis never participates in the matrix multiply. (B, D_in) @ (D_in, H) contracts over D_in and leaves B untouched. Each row of the output is computed from the corresponding row of the input and nothing else. Examples in a batch do not interact in a linear layer, which is why the same parameters work for a batch of 8 and a batch of 1600.
One orientation, stated explicitly
We wrote W1 as (in, out) and used x @ W1. PyTorch’s nn.Linear stores its weight the other way round, as (out, in), and applies x @ weight.T. Both conventions are correct; mixing them without noticing is how a (2, 16) tensor ends up somewhere expecting (16, 2).
We will verify the equivalence numerically at the end of the chapter rather than asserting it now. In the meantime, when you read a linear computation in unfamiliar code, the questions worth asking are always the same:
Which dimension of the weight is input width?
Which is output width?
Is the input multiplied on the left or the right?
Is there a transpose, and is it on the weight or the input?
What is the bias shaped, and which axis does it broadcast over?
Answer those five and you can read any linear layer in any codebase, including the ones an assistant produces at three in the morning with the orientation flipped.
Initialization decides which units are allowed to differ
We wrote torch.randn(D_in, H) / math.sqrt(D_in) without justifying it. This is the first chapter where the model has more than a couple of parameters, so the choice now has consequences worth deriving.
Why the scale is what it is
The pre-activation for hidden unit j is a sum over input features:
z1[b, j] = sum over i of x[b, i] * W1[i, j] + b1[j]
If the inputs have roughly unit variance and the weights are drawn independently with variance s², then summing D_in such products gives a variance of roughly D_in * s². Choosing s = 1 / sqrt(D_in) makes that product equal one, so the pre-activations arrive at the same scale as the inputs rather than growing or shrinking layer by layer. Measure it:
with torch.no_grad():
z1 = X_train @ W1 + b1
h = torch.relu(z1)
logits = h @ W2 + b2
print("z1 std", round(z1.std().item(), 4), " mean", round(z1.mean().item(), 4))
print("h std", round(h.std().item(), 4),
" fraction exactly zero", round((h == 0).float().mean().item(), 4))
print("logits std", round(logits.std().item(), 4))
z1 std 0.9377 mean 0.0047
h std 0.5687 fraction exactly zero 0.4946
logits std 0.9424
The first-layer prediction lands almost exactly on target: predicted 1.0, observed 0.938. About half the ReLU outputs are exactly zero, which is what you expect when the pre-activations are centered near zero and the function clips everything negative.
The second layer is a good reminder that these are approximations resting on assumptions. The variance argument assumes the terms being summed are roughly independent, and here all sixteen hidden units are functions of the same two inputs, so they are strongly correlated. The rule still put the logits at a sane scale, but do not expect the arithmetic to be exact when its premise is not.
The consequence of getting the scale badly wrong is now easy to anticipate. Scale the first-layer weights up dramatically and the pre-activations and surviving ReLU outputs grow with them; if later-layer weights are also oversized, the logits can become enormous and the model begins in a state of extreme, arbitrary confidence. log_softmax itself is designed to remain numerically stable, but the optimization problem can still begin in a poor regime.
Scale the weights down too far and the logits cluster near zero, so predictions begin close to uniform. In a tiny network optimization may recover from either choice. As depth increases, badly chosen scales can compound across layers and make training progressively harder or even fail.
PyTorch’s own layers use a fan-in-based rule of the same family. nn.Linear initializes its weight uniformly in ±1/sqrt(fan_in):
import torch.nn as nn
layer = nn.Linear(2, 16)
print("1/sqrt(fan_in) =", 1 / math.sqrt(2))
print("observed range:", layer.weight.min().item(), "to", layer.weight.max().item())
1/sqrt(fan_in) = 0.7071067811865475
observed range: -0.7010772228240967 to 0.6354132294654846
Both rules scale with fan-in, but they are not numerically identical. Our normal initialization uses standard deviation 1/sqrt(fan_in), while the default nn.Linear uniform bound produces a smaller variance. The common idea is that the scale depends on how many inputs feed each unit, not that the two distributions preserve exactly the same variance.
That is why the idea matters more than the recipe. When you later meet initializers such as kaiming_normal_ or xavier_uniform_, they will be variations on reasoning you have already encountered rather than names to memorize.
Why the randomness is there
Scale is the easier half. Here is the half that decides whether the hidden layer has sixteen units or effectively one.
Initialize every hidden unit identically. Give all sixteen the same incoming weights and the same outgoing weights, while keeping the two output classes distinguishable so that gradients are not trivially zero:
W1 = torch.full((D_in, H), 0.5, requires_grad=True)
b1 = torch.zeros(H, requires_grad=True)
W2 = torch.tensor([[0.25, -0.25]]).repeat(H, 1).requires_grad_()
b2 = torch.zeros(C, requires_grad=True)
Before training, take one backward pass and look at the structure of the gradients rather than their magnitudes:
loss = cross_entropy(forward(X_train), y_train)
loss.backward()
print("W1.grad columns all identical:", bool((W1.grad == W1.grad[:, :1]).all()))
print("W1.grad[:, :3] =", W1.grad[:, :3].tolist())
print("b1.grad[:3] =", b1.grad[:3].tolist())
print("W2.grad rows all identical: ", bool((W2.grad == W2.grad[:1]).all()))
W1.grad columns all identical: True
W1.grad[:, :3] = [[0.04355235, 0.04355235, 0.04355235],
[0.03529285, 0.03529285, 0.03529285]]
b1.grad[:3] = [0.10046272, 0.10046272, 0.10046272]
W2.grad rows all identical: True
This follows from the chain rule and could have been predicted without running anything. The gradient arriving at hidden unit j depends on the outgoing weights W2[j, :], which are identical across j. The gradient flowing into W1[:, j] depends on that incoming signal and on the input, which is shared. So every unit receives exactly the same gradient. Identical starting values plus identical gradients plus an identical update rule means identical values after the step, and the argument repeats at every step thereafter.
Train it for two thousand steps and check whether the prediction holds:
step= 0 loss=1.0926 train_acc=0.493 val_acc=0.512
step= 100 loss=0.6447 train_acc=0.606 val_acc=0.610
step= 500 loss=0.6436 train_acc=0.609 val_acc=0.600
step=2000 loss=0.6436 train_acc=0.611 val_acc=0.603
print("distinct columns of W1:", torch.unique(W1.detach(), dim=1).shape[1])
print("W1[:, :3] =", W1.detach()[:, :3].tolist())
with torch.no_grad():
h = torch.relu(X_train @ W1 + b1)
print("distinct hidden activation columns:", torch.unique(h, dim=1).shape[1])
distinct columns of W1: 1
W1[:, :3] = [[0.38603309, 0.38603309, 0.38603309],
[0.39404541, 0.39404541, 0.39404541]]
distinct hidden activation columns: 1
Sixteen hidden units, one distinct hidden function among them, bit-for-bit identical after two thousand updates. The parameters moved, the gradients were finite and non-zero throughout, and the loss decreased — but the nominal sixteen-unit layer has collapsed to the expressive diversity of a single hidden ReLU unit.
Its roughly 60% accuracy happens to lie near the affine reference on this dataset, but the model itself is not affine. The diagnostic evidence is the symmetry: sixteen nominal units are still computing one hidden feature.
Two qualifications keep this from becoming a slogan.
The symmetry has to hold on both sides to persist. Randomize W1 while leaving W2 constant and the units compute different functions immediately, so they receive different gradients, and the network trains normally to 99.5%. One source of asymmetry is enough.
And zero is not automatically fatal; it depends on where. Our biases start at exactly zero in every working version in this chapter. What breaks is initializing the weights to a constant. Set every weight and bias in this network to zero and the failure is total but entirely explicable:
step= 0 loss=0.6931 train_acc=0.493 val_acc=0.512
step= 100 loss=0.6931 train_acc=0.507 val_acc=0.488
step=2000 loss=0.6931 train_acc=0.507 val_acc=0.488
for name, p in zip(["W1", "b1", "W2", "b2"], parameters):
print(f"{name}.grad norm = {p.grad.norm().item():.6f}")
print("b2.grad =", b2.grad.tolist())
print("class 1 fraction in train:", y_train.float().mean().item())
W1.grad norm = 0.000000
b1.grad norm = 0.000000
W2.grad norm = 0.000000
b2.grad norm = 0.009723
b2.grad = [0.00687499949708581, -0.00687499949708581]
class 1 fraction in train: 0.5068749785423279
With W2 at zero, nothing downstream of the hidden layer can transmit a signal back, so W1, b1 and W2 all receive exactly zero. Only b2 gets anything, and its gradient is 0.5 − 0.50687 = −0.00687, which is precisely the difference between the uniform prediction and the class frequency. The model trains one thing: the base rate. Accuracy 0.507 on train and 0.488 on validation are the two class frequencies. Nothing is broken in the implementation. The initialization has placed the hidden part of the network in a symmetric state where its gradient signal is exactly zero. Gradient descent can still adjust b2, so the full parameter vector is not initially stationary; it simply learns the class base rate while the hidden representation remains frozen.
Once b2 reaches the best constant prediction, the remaining trainable directions available from this symmetric state provide no gradient signal capable of breaking the hidden units apart.
Initialization does not merely choose starting numbers. It can determine whether useful gradient signal reaches particular parameters and whether nominally distinct units begin with any opportunity to become different.
The objective: logits, log-softmax, and what cross entropy computes
The forward pass ends at (B, C). Each row holds two unnormalized scores. Those are logits: real numbers on any scale, where only differences between entries carry meaning, since adding a constant to a whole row leaves the softmax unchanged.
Turning logits into a loss for class-index targets goes through four steps:
logits (B, C) arbitrary real scores
↓ log_softmax over dim=1
log probabilities (B, C) each row sums to 1 after exponentiation
↓ select the entry at the target index
per-example score (B,) log probability assigned to the correct class
↓ negate and mean
loss scalar
Written out, with a comment on each line:
def cross_entropy(logits, targets):
log_probs = torch.log_softmax(logits, dim=1) # (B, C)
rows = torch.arange(targets.shape[0], device=logits.device)
picked = log_probs[rows, targets]
return -picked.mean() # scalar
The indexing line is worth reading slowly. torch.arange(B) is [0, 1, 2, ...] and targets is a (B,) tensor of class indices, so the pair selects element targets[b] from row b: advanced indexing along both axes at once. This hand-written version assumes class-index targets, so targets must contain integer class indices with shape (B,). PyTorch’s full cross_entropy API also supports probability targets in a different shape and dtype; that is a different contract from the one we are implementing here.
Verify it against the trusted implementation
The point of writing a mechanism by hand is not to use your version. It is to have something to compare against.
import torch.nn.functional as F
logits = forward(X_train)
mine = cross_entropy(logits, y_train)
theirs = F.cross_entropy(logits, y_train)
print("manual:", mine.item())
print("F.cross_entropy:", theirs.item())
print("allclose:", torch.allclose(mine, theirs))
manual: 0.8267903327941895
F.cross_entropy: 0.8267903327941895
allclose: True
Equal values are good evidence and incomplete evidence. The quantity training actually consumes is the gradient, and two formulas can agree on a value while differing in how they differentiate. Compare those too:
g_mine = torch.autograd.grad(cross_entropy(forward(X_train), y_train), parameters)
g_theirs = torch.autograd.grad(F.cross_entropy(forward(X_train), y_train), parameters)
for name, a, b in zip(["W1", "b1", "W2", "b2"], g_mine, g_theirs):
print(f"{name}: max abs difference = {(a - b).abs().max().item():.3e} "
f"allclose = {torch.allclose(a, b)}")
W1: max abs difference = 0.000e+00 allclose = True
b1: max abs difference = 0.000e+00 allclose = True
W2: max abs difference = 0.000e+00 allclose = True
b2: max abs difference = 0.000e+00 allclose = True
In this run the values and gradients agree exactly. The important claim is numerical equivalence, not guaranteed bit-for-bit identity across every device, kernel and PyTorch version, which is why allclose is the right general check. This pattern is worth making a habit:
When you reimplement a framework mechanism in order to understand it, derive a quantity both ways and compare numbers. Compare the value and the gradient, because those are separate claims.
Why log_softmax rather than the definition
The textbook formula is -log(p[target]) where p = exp(logits) / exp(logits).sum(). Implemented literally:
def naive_cross_entropy(logits, targets):
probs = logits.exp() / logits.exp().sum(dim=1, keepdim=True)
return -probs[torch.arange(targets.shape[0]), targets].log().mean()
big = torch.tensor([[100.0, 200.0], [1.0, 2.0]])
t = torch.tensor([1, 0])
print("naive: ", naive_cross_entropy(big, t).item())
print("stable:", cross_entropy(big, t).item())
print("torch: ", F.cross_entropy(big, t).item())
naive: nan
stable: 0.65663081407547
torch: 0.65663081407547
exp(200) overflows float32 to infinity, and the ratio of infinities is nan. The same quantity can be evaluated stably through the log-sum-exp identity: shift the logits by their maximum before exponentiating, then account for that shift algebraically. torch.log_softmax provides the stable fused operation, avoiding the explicit exp(200) that broke the naive implementation. Logits of 200 are not exotic; a slightly too-large initialization or a learning rate that briefly overshoots will produce them. Chapter 11 treats non-finite values systematically. Here it is enough to know that the stable spelling is not a stylistic preference.
Passing probabilities where logits belong
This is one of the most common mistakes in generated classifier code:
probs = torch.softmax(logits, dim=1)
loss = F.cross_entropy(probs, targets)
It is legal. The shapes are right, the dtypes are right, nothing warns. And what happens next is more interesting than “it breaks.”
F.cross_entropy applies log_softmax to whatever you hand it. Given probabilities, it applies softmax again. Look at one example:
logits = torch.tensor([[2.0, -1.0]])
t = torch.tensor([0])
probs = torch.softmax(logits, dim=1)
print("logits:", logits.tolist(), " probs:", [round(v, 4) for v in probs[0].tolist()])
print("cross_entropy(logits, t) =", F.cross_entropy(logits, t).item())
print("cross_entropy(probs, t) =", F.cross_entropy(probs, t).item())
logits: [[2.0, -1.0]] probs: [0.9526, 0.0474]
cross_entropy(logits, t) = 0.04858732968568802
cross_entropy(probs, t) = 0.33966848254203796
The model is quite confident and correct, which the proper loss reports as 0.049. The mistaken version reports 0.340. The reason is structural: probabilities live in [0, 1], so the gap between the two entries of a row can never exceed 1, no matter how certain the model becomes. Treated as logits, a maximum gap of 1 gives a minimum achievable loss of
-log(sigmoid(1)) = log(1 + e⁻¹) ≈ 0.3133
Check that the floor is real, using probabilities that are as confident as probabilities can get:
perfect = torch.tensor([[0.0, 1.0], [1.0, 0.0]])
print("CE of perfectly confident probabilities:", F.cross_entropy(perfect, torch.tensor([1, 0])).item())
print("CE of perfectly confident logits: ",
F.cross_entropy(torch.tensor([[-20., 20.], [20., -20.]]), torch.tensor([1, 0])).item())
CE of perfectly confident probabilities: 0.31326165795326233
CE of perfectly confident logits: 0.0
Now train the working network with the mistake in place, logging the reported loss alongside the loss that should have been reported:
step= 0 reported_loss=0.7157 loss_on_logits=0.8268 train_acc=0.281 val_acc=0.278
step= 100 reported_loss=0.4406 loss_on_logits=0.2341 train_acc=0.931 val_acc=0.942
step= 500 reported_loss=0.3582 loss_on_logits=0.0811 train_acc=0.993 val_acc=0.995
step=1000 reported_loss=0.3443 loss_on_logits=0.0561 train_acc=0.995 val_acc=0.995
step=2000 reported_loss=0.3347 loss_on_logits=0.0386 train_acc=0.996 val_acc=0.995
step=4000 reported_loss=0.3278 loss_on_logits=0.0263 train_acc=0.998 val_acc=0.995
The model still learns on this problem. Applying softmax preserves the ordering of the logits, so the predicted class from argmax is unchanged. The new objective also still rewards increasing the target class relative to the others, so optimization can improve classification accuracy.
But the extra softmax changes the scale and geometry of the gradients. In this experiment convergence is slower: at step 100 the correct objective has reached 0.1395 while the corresponding loss on the logits in the mistaken run is still 0.2341.
What is unambiguously destroyed is the loss as a measurement. It decreases monotonically toward 0.3278 and it will never go below 0.3133 however good the model becomes. Anyone reading that curve and concluding the model has stopped improving would be reading an artifact of the mistake. Anyone comparing that number against a published cross-entropy baseline would be comparing incompatible quantities.
So the honest statement is narrower and more useful than the usual warning:
F.cross_entropyexpects logits because it applieslog_softmaxinternally. Passing probabilities does not necessarily prevent learning; it silently changes the objective and makes the reported loss non-comparable and floor-limited. “The loss went down” is not evidence that the loss was formulated correctly.
If you want probabilities, compute them for inspection, outside the loss:
with torch.no_grad():
probabilities = torch.softmax(forward(X_val), dim=1)
Backward, and then the update, which is not backward
The gradient step is the part of this chapter Chapter 1 already taught, so it needs less explanation and one specific warning.
loss.backward()
That single call walks the graph from the scalar loss back through the second matrix multiply, the ReLU, the first matrix multiply and the two bias additions, accumulating into the four leaves. We wrote the forward computation; PyTorch derived the backward computation. The gradient of each parameter has the same shape as the parameter, which is worth confirming once because it is the fastest sanity check available:
W1.grad (2, 16) b1.grad (16,) W2.grad (16, 2) b2.grad (2,)
Then the update, in place, outside the graph:
with torch.no_grad():
for p in parameters:
p -= learning_rate * p.grad
torch.no_grad() is there because the optimizer’s parameter update is not part of the function we want autograd to differentiate. In normal grad mode, PyTorch refuses this particular in-place update on a leaf tensor that requires gradients.
Inside no_grad(), the subtraction is not recorded in the autograd graph, so the existing leaf tensor can be updated in place while remaining the trainable object that future forward passes use.
The update that runs and changes nothing
Chapter 1 met a rebinding update in a two-parameter model and promised to revisit it where it becomes genuinely dangerous. This is that place. With parameters held in a list and updated in a loop, write p = p - lr * p.grad instead of p -= lr * p.grad:
with torch.no_grad():
for p in parameters:
p = p - learning_rate * p.grad # rebinding, not in-place
step= 0 loss=0.8268 train_acc=0.281 val_acc=0.278
step= 100 loss=0.8268 train_acc=0.281 val_acc=0.278
step= 500 loss=0.8268 train_acc=0.281 val_acc=0.278
No exception, ever. In Chapter 1’s single-variable version the rebinding at least produced a RuntimeError on the following iteration, because the name w was rebound to a tensor that no longer required gradients. Here it cannot. p is a loop variable; rebinding it changes what p points at for the remainder of that iteration and nothing else. The list still holds the original four tensors, the next forward pass reads them, and they have never been touched. The new tensors are discarded immediately.
The loss is bit-identical across these steps, which is strong evidence that the state may not be changing and is worth checking immediately. It is not proof: a real update can be zero, lie in a flat direction, or be too small to change the reported loss at the displayed precision.
The decisive measurement is the parameter snapshot. Collect the evidence explicitly:
before = [p.detach().clone() for p in parameters]
# ... run the loop with the rebinding update ...
print("grad norms:", [round(p.grad.norm().item(), 4) for p in parameters])
print("max movement:", [(p.detach() - old).abs().max().item()
for p, old in zip(parameters, before)])
grad norms: [0.2437, 0.1834, 0.623, 0.231]
max movement: [0.0, 0.0, 0.0, 0.0]
Healthy, non-zero gradients on all four parameters. Zero movement on all four. This is the distinction the book keeps returning to, now in its sharpest form:
A populated
.gradis evidence about the backward pass. A changed parameter is evidence about the update. They are different claims, and establishing one does not establish the other.
The check that separates them costs two lines and belongs in your reflexes:
before = [p.detach().clone() for p in parameters]
with torch.no_grad():
for p in parameters:
p -= learning_rate * p.grad
for name, p, old in zip(["W1", "b1", "W2", "b2"], parameters, before):
observed = p.detach() - old
predicted = -learning_rate * p.grad
print(f"{name}: max|observed - predicted| = {(observed - predicted).abs().max().item():.3e} "
f"max|movement| = {observed.abs().max().item():.4f}")
W1: max|observed - predicted| = 5.751e-08 max|movement| = 0.0657
b1: max|observed - predicted| = 0.000e+00 max|movement| = 0.0637
W2: max|observed - predicted| = 2.608e-08 max|movement| = 0.1159
b2: max|observed - predicted| = 0.000e+00 max|movement| = 0.0817
This does more than confirm that something moved. It confirms that each parameter moved by exactly the amount the update rule specifies, to float32 precision. Under plain gradient descent the predicted movement is -lr * grad; under an optimizer with momentum or per-parameter scaling it will not be, and knowing which rule you are checking against is part of the check.
Finally, gradient clearing. Chapter 1 established the mechanism in detail — backward() adds to .grad rather than assigning — so it needs only a placement note here. In the loop we clear before backward():
for p in parameters:
if p.grad is not None:
p.grad.zero_()
loss.backward()
Gradients must be cleared before the next backward() whose result you want to interpret independently. Two common placements are immediately before backward(), as shown here, or immediately after the parameter update. Clearing them after backward() but before the update would erase the gradients the update needs.
The is not None guard exists because a fresh leaf has no .grad until the first backward pass populates it. The guard is also, as we saw with the non-leaf initialization, capable of hiding a parameter that never receives a gradient at all. It is a convenience, not a check.
The complete program
Everything assembled, with held-out evaluation:
import math
import torch
torch.manual_seed(0)
n = 2000
X = torch.randn(n, 2)
y = ((X[:, 0] > 0) ^ (X[:, 1] > 0)).long()
perm = torch.randperm(n)
n_train = int(0.8 * n)
X_train, y_train = X[perm[:n_train]], y[perm[:n_train]]
X_val, y_val = X[perm[n_train:]], y[perm[n_train:]]
D_in, H, C = 2, 16, 2
torch.manual_seed(1)
W1 = (torch.randn(D_in, H) / math.sqrt(D_in)).requires_grad_()
b1 = torch.zeros(H, requires_grad=True)
W2 = (torch.randn(H, C) / math.sqrt(H)).requires_grad_()
b2 = torch.zeros(C, requires_grad=True)
parameters = [W1, b1, W2, b2]
def forward(x):
z1 = x @ W1 + b1
h = torch.relu(z1)
return h @ W2 + b2
def cross_entropy(logits, targets):
log_probs = torch.log_softmax(logits, dim=1)
picked = log_probs[torch.arange(targets.shape[0]), targets]
return -picked.mean()
def accuracy(logits, targets):
return (logits.argmax(dim=1) == targets).float().mean().item()
learning_rate = 0.5
for step in range(2001):
logits = forward(X_train)
loss = cross_entropy(logits, y_train)
if step % 500 == 0:
with torch.no_grad():
val_logits = forward(X_val)
print(f"step={step:4d} loss={loss.item():.4f} "
f"train_acc={accuracy(logits, y_train):.3f} "
f"val_loss={cross_entropy(val_logits, y_val).item():.4f} "
f"val_acc={accuracy(val_logits, y_val):.3f}")
for p in parameters:
if p.grad is not None:
p.grad.zero_()
loss.backward()
with torch.no_grad():
for p in parameters:
p -= learning_rate * p.grad
step= 0 loss=0.8268 train_acc=0.281 val_loss=0.8097 val_acc=0.278
step= 500 loss=0.0589 train_acc=0.994 val_loss=0.0511 val_acc=0.995
step=1000 loss=0.0408 train_acc=0.996 val_loss=0.0335 val_acc=0.995
step=1500 loss=0.0327 train_acc=0.997 val_loss=0.0262 val_acc=0.995
step=2000 loss=0.0279 train_acc=0.998 val_loss=0.0221 val_acc=0.995
Forty lines, no framework abstractions, and 99.5% on four hundred examples the update step never touched.
The validation split is doing real work here, and it is doing a different job from the training loss. Chapter 1 showed a model driving its training loss to exactly zero while being wrong everywhere it had not been shown. The two columns answer two questions:
falling training loss → the optimizer is reducing the objective it was given
held-out accuracy → the model behaves correctly on data it was not fitted to
In this run both agree, which is the uninteresting case. The interesting cases are the ones where they diverge, and you cannot notice a divergence you are not measuring. Chapter 14 turns held-out evaluation into a proper discipline; a single accuracy number is enough for now.
One deliberate omission: this trains on the full batch every step, so the gradient is the exact gradient of the objective over the training set. Real training uses mini-batches, which makes each step a noisy estimate and changes the character of the optimization. The mechanism is unchanged — same forward, same backward, same update — and Chapter 6 builds the data pipeline that supplies those batches.
Six questions, and the evidence that answers each
This chapter has produced five distinct failures. Line them up and notice how little they have in common except that the program ran.
| What was wrong | Loss behavior | What the tensors said |
|---|---|---|
| No activation | fell, then plateaued at 0.6928 | all leaves, gradients finite, parameters moving |
Non-leaf W1 |
fell to 0.29, plausible accuracy | W1.grad is None, W1 never changed |
| Symmetric init | fell, plateaued at 0.6436 | gradients identical across hidden units |
| Zero init | flat at 0.6931 | every gradient zero except b2 |
| Rebinding update | bit-identical every step | gradients healthy, movement exactly zero |
No single check catches all five, which is why “check the gradients” is not a debugging strategy. What generalizes is knowing which question a piece of evidence answers. There are six questions, and the previous three chapters gave us tools for only three of them.
Representation. Can this architecture express the function at all? Evidence can come from an algebraic argument about the model class, a known representational bound, a falsifying probe such as the affine-additivity residual, or comparison with a deliberately simpler baseline. The opening failure turned on exactly this question: healthy tensors and gradients do not by themselves establish that the architecture can represent the required function.
Computation. Is the forward pass computing the architecture you think it is? Evidence: shape traces, assertions on intermediate tensors, and comparison against an independent implementation of the same function.
Objective. Does the loss encode what you want, and is it consuming the kind of value it expects? Evidence: comparison against a trusted implementation on both value and gradient, and checking the range and meaning of what you pass in. Logits or probabilities, class indices or one-hot, sum or mean.
Differentiation. Does every parameter you intend to train have a path to the loss? Evidence: requires_grad, is_leaf, grad_fn, .grad, and the reachability check from Chapter 3.
Update. Do those parameters actually change? Evidence: snapshot before, compare after, and check the movement against what the update rule predicts.
Evaluation. Are you measuring the behavior you care about? Evidence: held-out performance and an appropriate baseline whose behavior and limitations you understand.
On this small problem, representation and objective are unusually cheap to check, so they are sensible places to start. In larger systems the cost of each check will vary. What matters more than a universal ordering is knowing which claim each piece of evidence actually supports. But the real value is not the ordering. It is being able to say, on seeing a piece of evidence, which of the six claims it supports and which it leaves entirely open.
“The gradients look fine” answers exactly one of these six questions.
The same network, spelled in PyTorch
Now the reveal. Here is our model in idiomatic PyTorch:
import torch.nn as nn
model = nn.Sequential(
nn.Linear(2, 16),
nn.ReLU(),
nn.Linear(16, 2),
)
That is where most tutorials begin, and it is exactly what we have already built. Prove it rather than asserting it. Copy our parameters into the module, remembering the orientation difference, and compare both the outputs and the gradients:
model = nn.Sequential(nn.Linear(2, H), nn.ReLU(), nn.Linear(H, C))
with torch.no_grad():
model[0].weight.copy_(W1.T)
model[0].bias.copy_(b1)
model[2].weight.copy_(W2.T)
model[2].bias.copy_(b2)
manual_logits = torch.relu(X_train @ W1 + b1) @ W2 + b2
module_logits = model(X_train)
print("logits max abs difference:", (manual_logits - module_logits).abs().max().item())
loss_manual = F.cross_entropy(manual_logits, y_train)
loss_module = F.cross_entropy(module_logits, y_train)
g_manual = torch.autograd.grad(loss_manual, [W1, b1, W2, b2])
g_module = torch.autograd.grad(
loss_module,
[model[0].weight, model[0].bias, model[2].weight, model[2].bias],
)
print("dL/dW1 vs dL/d(fc1.weight).T:", (g_manual[0] - g_module[0].T).abs().max().item())
print("dL/db1 vs dL/d(fc1.bias) :", (g_manual[1] - g_module[1]).abs().max().item())
print("dL/dW2 vs dL/d(fc2.weight).T:", (g_manual[2] - g_module[2].T).abs().max().item())
print("dL/db2 vs dL/d(fc2.bias) :", (g_manual[3] - g_module[3]).abs().max().item())
logits max abs difference: 0.0
dL/dW1 vs dL/d(fc1.weight).T: 0.0
dL/db1 vs dL/d(fc1.bias) : 0.0
dL/dW2 vs dL/d(fc2.weight).T: 0.0
dL/db2 vs dL/d(fc2.bias) : 0.0
Not close. Identical, in every element, for outputs and for all four gradients. Whatever nn.Linear and nn.Sequential are, they are not computing anything we were not already computing.
The one substantive difference is the transpose, and now is the moment to confirm it directly:
layer = nn.Linear(2, 16)
xb = torch.randn(8, 2)
print("weight shape:", tuple(layer.weight.shape))
print("layer(x) == x @ weight.T + bias:",
torch.allclose(layer(xb), xb @ layer.weight.T + layer.bias))
weight shape: (16, 2)
layer(x) == x @ weight.T + bias: True
nn.Linear(in, out) stores its weight as (out, in) and computes x @ weight.T + bias. We stored (in, out) and computed x @ W + b. Same function, different storage convention, and now you know which one you are looking at when a shape does not match your expectation.
The optimizer is the second piece of packaging:
optimizer = torch.optim.SGD(model.parameters(), lr=0.5)
optimizer.zero_grad()
logits = model(X_train)
loss = F.cross_entropy(logits, y_train)
loss.backward()
optimizer.step()
loss.backward() is unchanged and still does the differentiation. optimizer.step() reads the .grad fields that backward() populated and applies an update rule to the parameters it was constructed with. Those are the same two separate operations we have been writing by hand, which is why the distinction between “gradients exist” and “parameters moved” survives the move to the framework. It becomes, if anything, easier to get wrong, because an optimizer constructed with the wrong parameter list will silently update nothing while backward() continues populating gradients perfectly.
Line the two spellings up:
MANUAL PYTORCH
W1, b1 nn.Linear(2, 16) weight stored as (out, in)
torch.relu(...) nn.ReLU() / F.relu
W2, b2 nn.Linear(16, 2)
def forward(x): ... Module.forward(self, x)
parameters = [W1, b1, W2, b2] model.parameters()
p -= lr * p.grad optimizer.step()
clear/reset parameter gradients optimizer.zero_grad()
manual cross_entropy(...) F.cross_entropy(...)
Every entry on the right is a name for a mechanism on the left. Nothing in the right-hand column computes anything the left-hand column did not.
What the right-hand column does buy is on the other side of the arrow: the list on the left has to be maintained by hand. Add a third layer and you write two more tensors, add them to parameters, add them to the update loop, add them to the device move, add them to the save dictionary, and every one of those is a place to forget one. model.parameters() removes the need to maintain that list manually — provided the parameters and submodules have been registered correctly. The framework can reliably walk registered structure; Chapter 5 is about what counts as registered structure and how seemingly ordinary Python containers can fall outside it.
How it knows is Chapter 5, and it is a better question than it sounds. Assigning a tensor to self.something does not make it a parameter. Assigning modules to a plain Python list does not register them. The rules for what gets discovered by parameters(), state_dict(), .to(device) and .eval() all follow from a single organizing idea, and a particularly nasty category of PyTorch bug is a parameter that exists, participates in computation, and receives gradients while remaining invisible to some of the framework machinery that was expected to manage it.
You are now in a good position to understand that failure, because you have built the thing being registered.
Using AI on a network you did not write
The failures in this chapter are also easy for an assistant to produce or overlook, for the same reason they are easy for a human to miss: the code looks plausible. A missing nn.ReLU() between two nn.Linear calls reads as an ordinary two-layer network. A torch.softmax before F.cross_entropy reads as careful. Ask “is this correct?” and you will often get a fluent argument that it is.
The move that works is the one from the previous three chapters: make it commit to checkable claims about the specific program before it is allowed to change anything.
Here is a small PyTorch network built from raw tensors. Do not rewrite it and
do not propose a fix yet.
First, for every trainable tensor:
- its symbolic shape, and what each axis indexes;
- the path from that tensor to the loss;
- whether .grad should be populated after backward(), and why;
- which statement in the update should change it.
Then answer these four separately:
1. What is the set of functions this architecture can express? Specifically,
is the composed forward pass affine in the input, and how would I test that
empirically without reading the source?
2. Does the forward computation match the architecture the code appears to
intend?
3. What does the loss expect as input — logits, probabilities, class indices,
one-hot targets — and what is it actually being given?
4. Give me three measurements that would distinguish "this architecture cannot
represent the task" from "the training loop is broken". For each, say what
result would support which conclusion.
Give me the code for those measurements. I will run them and tell you what I see.
The last line is the one that changes the interaction. An assistant asked to fix code will produce a fix; an assistant asked to produce a discriminating experiment will produce something you can run, and running it either confirms its model of your program or shows you exactly where its model is wrong. Both outcomes are worth more than a rewrite.
Question 1 is deliberately the first. It is the question that the tooling from Chapters 1 to 3 cannot answer, it is answerable in two lines of arithmetic, and it eliminates or confirms an entire branch of explanation before anyone touches a hyperparameter.
Ask an assistant to characterize the system before asking it to replace the system.
What you should now be able to answer
Here is a network of the kind you will be handed. It runs. Work through it before reading the answers.
import torch
import torch.nn.functional as F
W1 = torch.randn(4, 32, requires_grad=True) * 0.05
b1 = torch.zeros(32, requires_grad=True)
W2 = (torch.randn(32, 3) * 0.05).requires_grad_()
b2 = torch.zeros(3, requires_grad=True)
params = [W1, b1, W2, b2]
def forward(x):
h = torch.relu(x @ W1 + b1)
return h @ W2 + b2
def loss_fn(logits, targets):
return F.cross_entropy(torch.softmax(logits, dim=1), targets)
for step in range(1000):
loss = loss_fn(forward(X), y)
for p in params:
if p.grad is not None:
p.grad.zero_()
loss.backward(retain_graph=True)
with torch.no_grad():
for p in params:
if p.grad is not None:
p -= 0.01 * p.grad
What are the model’s trainable values, and how many numbers is that? W1, b1, W2 and b2: 4×32 + 32 + 32×3 + 3 = 259 numbers. The input has four features, the hidden layer has thirty-two units, and there are three classes.
Why is each tensor shaped the way it is? W1 is (in, out) because the code multiplies as x @ W1, so axis 0 must match the input width and axis 1 defines the hidden width. b1 is (32,), one offset per hidden unit, broadcast across the batch. W2 maps thirty-two hidden units to three classes. b2 is one offset per class. Note that nn.Linear would store these transposed.
Which tensors are leaves? b1, W2 and b2. W1 is not: requires_grad=True was applied before the multiplication by 0.05, so W1 is the output of a recorded MulBackward0. Check with W1.is_leaf.
Which .grad fields will populate? b1.grad, W2.grad and b2.grad. W1.grad stays None, because non-leaves do not retain gradients. Reading it emits a UserWarning.
Which parameters will actually change? Only those three. W1 holds its initial random values for all thousand steps, and the if p.grad is not None guard skips it without comment. The model trains a classifier on top of thirty-two frozen random features.
What is retain_graph=True doing there? Papering over the previous answer. Because W1 is a non-leaf created outside the loop, its MulBackward0 node is shared by every iteration’s graph, and without retain_graph the second backward() raises. The flag removes the exception; it does not make W1 trainable.
Can this architecture represent a nonlinear boundary? Yes. There is a ReLU between the two affine maps, so the composition is not affine. Confirm with the additivity residual rather than by reading the code.
Does the loss receive what it expects? No. F.cross_entropy applies log_softmax internally, so passing torch.softmax(logits, dim=1) softmaxes twice. Training may still proceed, but the reported loss is floor-limited and not comparable with any other cross-entropy number. The floor depends on the number of classes. With C classes, the limiting case as the predicted probability vector approaches one-hot gives a floor of log(1 + (C-1)/e), which is about 0.313 for two classes and 0.551 for three.
What would you check first? Not the learning rate. In order: is the composed forward affine (representation), does the loss consume logits (objective), does every intended parameter have a populated .grad (differentiation), and did every intended parameter move (update). Three of the four are wrong here, and none of them are hyperparameters.
How would you verify a repair? Name the quantity and observe it. After creating W1 as (torch.randn(4, 32) * 0.05).requires_grad_(): W1.is_leaf is True, W1.grad is populated after backward(), a before/after snapshot shows W1 moving by -lr * W1.grad, and retain_graph=True can be removed without an exception. After removing the softmax, the reported loss is no longer constrained by the probability-as-logits floor. For this three-class example that limiting floor is about 0.551, not 0.313. Then verify the repaired objective against the hand-written log_softmax implementation on both value and gradient.
Exercises
These are written to be run, and they map onto the notebook that accompanies this chapter.
-
Prove the collapse both ways. For the network without an activation, compute
W_eq = W1 @ W2andb_eq = b1 @ W2 + b2and confirm the logits match. Then extend to three linear layers and derive the equivalent single affine map before checking it. Finally, runaffine_residualon all three networks and on the version with ReLU, and explain which test you would rather have when you cannot see the parameters. -
Build an independent linear reference. Use the dense grid search to find the best linear boundary it discovers on the XOR training data, then train a single-layer model (
X @ W + bonly) and see how close optimization gets to that reference.
Keep the boundary selected from the training data and evaluate that same boundary on validation data rather than searching the validation labels separately. Then change the label rule to y = (X[:, 0] > 0).long(), which is linearly separable, and repeat.
Explain which observations concern optimization, which concern representation, and why the grid search is an empirical reference rather than a proof of the exact optimum. Now change the label rule to y = (X[:, 0] > 0).long(), which is linearly separable, and repeat both measurements. Report all four numbers and say which pair tells you something about your training code and which pair tells you something about the problem. Then rotate the XOR inputs by 45 degrees, predict whether the ceiling moves, and check.
-
Count the folds. Train the ReLU network with
H = 1, 2, 3, 4, 8, 16, recording validation accuracy for each. Predict, before running, the smallest hidden size that could in principle separate the four quadrants. Then run each size across five seeds and report the spread rather than a single number — with a fixed seed the results are not monotonic inH, and working out why that is not a contradiction is the real exercise. Finally, plot the decision regions of the best model and count the linear pieces. -
Make the manual loss disagree. Write a version of
cross_entropythat usessum()instead ofmean(), and one that indexes with a one-hot multiply instead of advanced indexing. For each, compare value and gradient againstF.cross_entropyand explain any difference. Then train with thesum()version at the same learning rate and account for what happens. -
Locate the floor. Derive the minimum achievable value of
F.cross_entropy(probs, targets)forC = 2,C = 3andC = 10classes, then confirm each empirically with maximally confident probability vectors. Explain why the floor depends on the number of classes. -
Break the symmetry with one number. Start from the fully symmetric initialization and perturb exactly one element associated with hidden unit 0. Do this once in
W1and once inW2.
Before training, predict how many symmetry groups of hidden units now exist. Then take one backward pass and identify which gradients remain identical and which diverge. Train both versions and verify whether the original group of fifteen units remains internally symmetric.
The point is not that perturbing W1 and W2 necessarily produces different group counts; it is to trace how a single asymmetry propagates through the forward and backward computations.
- Two failures, almost the same symptom. Compare the network with no activation against the fully zero-initialized nonlinear network. Both settle near
log 2, but for completely different reasons.
At an early training step, find the smallest measurement that separates them. In the affine network the gradients are valid and the architecture is the limiting factor. In the zero-initialized network the hidden-layer gradient signal is exactly zero. Explain why the same-looking loss curve supports two different diagnoses. Write the single smallest check that distinguishes them, and argue for whether it belongs in the training loop permanently.
-
Verify an unfamiliar layer. Pick a layer you have not used —
nn.Bilinear,nn.Embedding, ornn.Conv1d— read its documented shape contract, implement the forward computation with raw tensor operations, and compare outputs and gradients against the real layer on random input. Report the largest discrepancy you find and explain it. -
Audit a generated network. Ask an assistant for a small classifier written with raw tensors, then run the full six-question audit on it before running the code: representation, computation, objective, differentiation, update, evaluation. Record which questions you could answer by reading and which required a measurement.
Next: what the framework is organizing
We now have a complete learning system with no framework abstractions in it. Four tensors are the model. A function combines them with an input. A loss reduces the output to a scalar. backward() differentiates the recorded computation. An in-place update modifies the leaves. Everything in that list is a mechanism from Chapters 1 through 3, at slightly larger scale.
We also have something the earlier chapters could not provide: a way to distinguish an implementation that is broken from an implementation that is correct and incapable. That distinction is what separates useful debugging from an afternoon of learning-rate roulette, and it becomes more valuable as models get larger, because larger models offer more places for both kinds of problem to hide.
What we do not have is a way to manage this at scale. Four parameters in a hand-maintained list is fine. Four hundred is not, and the failure modes are not interesting ones: a tensor left out of the update loop, a device move that missed one, a checkpoint that saved three of four weights. Those are bookkeeping errors, and bookkeeping is exactly what a framework should absorb.
The next chapter introduces nn.Module, and it is not new machinery. It is a way of organizing tensors, leaves and gradients that we have already built by hand, resting on one idea: objects contain other objects, and PyTorch knows how to walk the resulting structure. Every question that idea answers is a question we have just met in its manual form. Why does model.parameters() find a weight buried four levels deep? Why does model.to("cuda") move tensors you never named? Why does a module stored in a plain Python list silently vanish from state_dict()?
You already know what the network is. Next we find out what the framework is doing to it.
Recursive composition.
