Stacking layers

In chapter 5, you trained a single neuron and watched it slide a line between two clouds of points. The line was the only thing it could do. Try the same trainer on data that isn't linearly separable — say, points arranged in a checkerboard pattern, or a circle inside another circle — and the neuron fails. There's no single line that separates them.

The fix is almost embarrassingly simple: stack neurons. A layer is several neurons running in parallel on the same input. They each carve out their own line. Combine their outputs and you get curves, kinks, regions — anything you want, given enough neurons. That's the whole story of why deep learning works.

We're going to inspect a 2-layer multi-layer perceptron (MLP), train it on XOR, and watch it succeed where the single neuron failed. Three runnable cells in the browser, then a local forward pass you will reuse when the transformer block arrives.

1. The linear layer

A linear layer is y = Wx + b, where W is a matrix and b a bias vector. If x has length n and the layer has out neurons, then W is shape [out, n] and b is length out. The output is length out.

Conceptually, each row of W is the weight vector of one neuron. Each row of Wx + b is that neuron's pre-activation.

That's the workhorse of every dense neural network. Read it once, get used to the shapes, then reuse it everywhere.

2. The MLP forward pass

Now we chain two linear layers with a non-linearity in between. The non-linearity is critical: without it, two stacked linear layers collapse to a single linear layer (matrix multiplication is associative). The non-linearity in modern networks is almost always ReLU: relu(x) = max(0, x).

The full forward pass for our 2-input, 1-output network is:

Write this as a function. The chapter compares your output against a reference implementation on four test points (one per quadrant of the plane).

If your bars match the reference's, your forward pass is correct. The two should agree exactly because we use the same weights and the same arithmetic. Untrained, the network's predictions are mostly noise — that's expected. Training is where the layers earn their keep.

3. Train on XOR

XOR is the textbook case where a single neuron fails. Two clusters in opposite quadrants are class 0; the other two quadrants are class 1. No line through the plane separates them.

Train an MLP with hidden size 8 on a noisy XOR dataset. The chapter provides initMlp, mlpStep (full-batch GD, returns updated net + loss), and the dataset; the cell runs the loop.

The left panel shows the network's decision regions — every pixel is colored by what the trained MLP thinks the class probability is at that location. Warm = class 1, cool = class 0, faint = uncertain. After training, you should see a checkerboard-shaped pattern carving the plane into four regions, with the data points landing in matching colors.

The loss curve on the right should decrease, possibly with a couple of plateaus. MLP training is famously not always smooth — the first few iterations often go nowhere while the random init untangles itself, then loss drops fast.

If training fails (the regions are wrong, the loss plateaus high), try increasing hiddenSize, increasing iterations, or adjusting lr. Pretty much every neural-network failure mode has the same fix: more capacity, more time, or more careful step sizes.

4. Extend llm/nn.py

Add these helpers below the neuron code:

# [1]
Vector = list[float]
Matrix = list[Vector]
 
 
# [2]
def relu(x: Vector) -> Vector:
    return [max(0.0, value) for value in x]
 
 
# [3]
def linear(x: Vector, weight: Matrix, bias: Vector) -> Vector:
    return [
        sum(xi * wi for xi, wi in zip(x, row)) + b
        for row, b in zip(weight, bias)
    ]
 
 
def mlp_forward(
    x: Vector,
    w1: Matrix,
    b1: Vector,
    w2: Matrix,
    b2: Vector,
) -> Vector:
    # [4]
    hidden = relu(linear(x, w1, b1))
    # [5]
    return linear(hidden, w2, b2)

Read the shapes first, then the code:

• [1] gives names to the shapes: a vector is one row, a matrix is many rows.
• [2] relu keeps positive signals and clips negative ones to zero. That is the non-linear bend the previous chapter did not have.
• [3] linear takes one input vector and many rows of weights. Each row is one neuron, and the list comprehension returns every neuron output as a new vector.
• [4] creates learned hidden features from the raw input.
• [5] mixes those hidden features into the final output.

The code is plain lists for now. That is deliberate: you can see every shape. In chapter 12, linear becomes torch.nn.Linear, but the contract stays the same.

What this fixes (and what it costs)

Fixes: any boundary you can imagine. Universal approximation guarantees that an MLP with enough hidden neurons can approximate any continuous function on a compact domain to arbitrary precision. You can fit XOR, spirals, concentric circles, anything.

Costs:

• Many more parameters. Hidden size 8 with 2 inputs is 8×2 + 8 + 8 + 1 = 33 parameters. Real networks have millions or billions.
• Local minima. Gradient descent on a non-convex loss can get stuck. Modern MLPs almost always escape because their loss landscapes are forgiving in high dimensions, but the theory is still mostly empirical.
• Hyperparameters. Hidden size, learning rate, initialization, optimizer choice — all of these affect whether training works. Chapter 7 starts that conversation.

Going further

• Karpathy's "Building micrograd" — derives the chain rule visually and builds an autograd engine that handles the gradient computation for arbitrary expressions.
• 3Blue1Brown's "What is backpropagation, really doing?" — gentle visual explanation of the gradient flow.
• Tinker (TensorFlow Playground) — drag a network's hyperparameters around and watch it train on built-in datasets. Good for building intuition.
• The full reference implementation lives in lib/ml/nn/.

Next up: gradient descent live — we've been calling mlpStep as a black box. Time to look inside, see why vanilla GD often gets stuck, and meet momentum and Adam.