Implementing a Network

A network is layers composed in sequence, where each layer's output becomes the next layer's input. The composition is what buys the extra expressiveness: a single linear layer can only draw one straight boundary, and no amount of training fixes that. Chaining a nonlinear layer in front of it can fit boundaries that are not lines at all.

The idea

Build a network by chaining functions. A two-layer model is f(x) = W₂ · ReLU(W₁x + b₁) + b₂, evaluated left to right: hidden activations, then output logits.

• Shape chaining. Each layer's out_features must equal the next layer's in_features, or the matmul fails.
• Parameter counting. 2 → 4 → 1 holds 2·4 + 4 = 12 weights and biases in the hidden layer, 4·1 + 1 = 5 in the output layer: 17 trainable numbers.
• Nonlinearity between layers. Remove the ReLU and W₂W₁x + b collapses into one linear map, which has exactly the representational power of a single layer.
• The training loop contract. Forward pass, compute loss, zero_grad(), backward(), step(). Skip zero_grad() and gradients accumulate across steps; call step() before backward() and the optimizer moves on stale gradients.

XOR is the standard test of this machinery. Its four examples are not linearly separable: no single line splits (0,1) and (1,0) from (0,0) and (1,1).

Worked example

Train two models on the XOR truth table, targets (0, 1, 1, 0).

• A single Linear(2, 1) plateaus at a loss of 0.6931 with every input predicted at probability 0.5. That is ln 2, the loss of guessing blindly.
• A Linear(2, 4) → ReLU → Linear(4, 1) network reaches 0.0008 and predictions (0.000, 1.000, 0.998, 0.000).

The MLP loss trace at learning rate 0.5 is 0.6877 at step 0, 0.1044 at step 100, 0.0083 at step 500, and 0.0008 at step 4000. Random initialization still matters: with torch.manual_seed(2) the same network stalls at 0.6931, because its ReLU units start below zero for every input and receive no gradient.

In code

import torch

X = torch.tensor([[0., 0.], [0., 1.], [1., 0.], [1., 1.]])
y = torch.tensor([[0.], [1.], [1.], [0.]])
lossf = torch.nn.BCEWithLogitsLoss()

torch.manual_seed(0)
model = torch.nn.Sequential(
    torch.nn.Linear(2, 4), torch.nn.ReLU(), torch.nn.Linear(4, 1)
)
opt = torch.optim.SGD(model.parameters(), lr=0.5)

step = 0
for _ in range(4001):
    opt.zero_grad()
    loss = lossf(model(X), y)
    loss.backward()
    opt.step()
    if step in (0, 100, 500, 4000):
        print(step, round(loss.item(), 4))       # 0.6877 0.1044 0.0083 0.0008
    step += 1

print([round(v, 3) for v in model(X).sigmoid().flatten().tolist()])   # [0.0, 1.0, 0.998, 0.0]
print(sum(p.numel() for p in model.parameters()))   # 17

Check yourself

1. Why can no Linear(2, 1) model fit XOR, no matter how long it trains?
2. What does opt.zero_grad() clear, and what goes wrong within a minute of training without it?
3. If you delete the ReLU from this network, what does the model reduce to, and what loss will it reach on XOR?

Key takeaways

• A network is a chain of layers, and only the nonlinearity between them makes depth worth paying for.
• The training loop is forward, loss, zero_grad, backward, step — in that order, every step.
• Small networks can fail from initialization as easily as from architecture; a stalled loss is a clue.