Chapter 14

Neural Networks from Scratch

A multilayer perceptron derived and built by hand: forward pass, backpropagation, initialization, and overfitting.

Run the companion notebook Download notebook AsciiDoc source

A neural network is a stack of affine maps and nonlinear functions whose parameters are chosen by gradient descent. For LLMs, the same idea appears inside every projection matrix, feed-forward block, and classifier head; only the tensors get larger. This chapter builds the small 5 → 10 → 4 multilayer perceptron from the companion notebook, derives its gradients, and trains a tiny reproducible NumPy version with plain SGD. The companion notebook runs the longer Adam training and overfitting experiment in JupyterLite.

14.1 Forward pass

With rows as examples, a one-hidden-layer MLP maps $\mX \,\in\, \R^{B \times 5}$ to four class scores by two affine layers and a ReLU:

\begin{aligned} \mZ_1 &= \mX\mW_1 + \vb_1, & \mH &= \max(0, \mZ_1), \\ \mZ_2 &= \mH\mW_2 + \vb_2, & \mP &= \softmax(\mZ_2). \end{aligned}\tag{14.1}

Here $\mW_1 \in \R^{5 \times 10}$ , $\vb_1 \in \R^{10}$ , $\mW_2 \in \R^{10 \times 4}$ , and $\vb_2 \in \R^4$ . Biases broadcast across rows, so the chapter network has 104 trainable scalars. ReLU keeps positive coordinates and sets the rest to zero, giving the model a piecewise-linear decision boundary instead of one linear classifier.

The forward pass should also decide what to cache. Backward needs the original inputs, the pre-ReLU activations, the hidden activations, and the final probabilities. It does not need to remember temporary shifted logits used only for numerical stability. Keeping the cache small matters in large networks, where activations can dominate memory during training.

Shapes are the fastest error check. If $\mX$ is $B \times 5$ , then $\mX\mW_1$ must be $B \times 10$ , and every gradient with respect to a parameter must have that parameter’s shape. Most NumPy bugs in hand-written networks are axis mistakes, not calculus mistakes. Writing shapes beside the code also makes broadcasting intentional instead of accidental, especially for bias vectors shared by every single example in the current minibatch.

The numerically stable softmax subtracts each row maximum before exponentiating. For labels $y_i$ , the batch mean cross-entropy is

L = -\frac{1}{B} \sum_{i=1}^{B} \log P_{i,y_i} .\tag{14.2}

Listing 14.1 Initialization for the notebook architecture

def initialize(seed=0, input_dim=5, hidden_dim=10, output_dim=4, dtype=np.float32):
    rng = np.random.default_rng(seed)
    return {
        "W1": (rng.normal(size=(input_dim, hidden_dim))
               * np.sqrt(2.0 / input_dim)).astype(dtype),
        "b1": np.zeros(hidden_dim, dtype=dtype),
        "W2": (rng.normal(size=(hidden_dim, output_dim))
               * np.sqrt(1.0 / hidden_dim)).astype(dtype),
        "b2": np.zeros(output_dim, dtype=dtype),
    }

14.2 Backpropagation in matrix form

Backpropagation is bookkeeping for the chain rule [rumelhart1986]. For an affine layer $\mZ = \mA\mW + \vb$ with upstream gradient $\bar{\mZ}$ , perturbing it gives $\dd L = \mathrm{tr}(\bar{\mZ}^{\T}\dd\mA\mW) + \mathrm{tr}(\bar{\mZ}^{\T}\mA\dd\mW) + \bar{\mZ}:\dd\vb$ . Matching coefficients yields

\bar{\mA} = \bar{\mZ}\mW^\T,\quad \bar{\mW} = \mA^\T\bar{\mZ},\quad \bar{\vb} = \sum_i \bar{\mZ}_i .\tag{14.3}

The ReLU gradient is just a mask: $\bar{\mZ}_1 = \bar{\mH} \odot \mathbf{1}[\mZ_1 > 0]$ . For softmax followed by cross-entropy, the Jacobian terms cancel. For one row with one-hot $\vy$ ,

\frac{\partial \ell}{\partial z_j} = \sum_k -y_k(\delta_{kj} - p_j) = p_j - y_j .\tag{14.4}

For the batch mean, start from $\bar{\mZ}_2 = (\mP - \mY)/B$ and apply (14.3) backward through the output affine, ReLU, and input affine. The tests check both the parameter gradient and the softmax-cross-entropy gradient against central finite differences, the method from Section B.8.

The order matters. Compute every gradient from the old parameters before applying any update; otherwise later gradients would mix old and new weights. Also average exactly once. If the loss is a mean over examples, the factor $1/B$ belongs in $\bar{\mZ}_2$ ; the affine and ReLU rules then propagate that scale automatically.

Listing 14.2 Forward and backward passes

def softmax_cross_entropy(logits, labels):
    shifted = logits - logits.max(axis=1, keepdims=True)
    log_probs = shifted - np.log(np.exp(shifted).sum(axis=1, keepdims=True))
    probabilities = np.exp(log_probs)
    loss = -log_probs[np.arange(len(labels)), labels].mean()
    return loss, probabilities


def forward(params, inputs, labels):
    z1 = inputs @ params["W1"] + params["b1"]
    hidden = np.maximum(z1, 0)
    logits = hidden @ params["W2"] + params["b2"]
    loss, probabilities = softmax_cross_entropy(logits, labels)
    return loss, {"X": inputs, "Z1": z1, "H": hidden, "P": probabilities}


def backward(params, cache, labels):
    grad_logits = cache["P"].copy()
    grad_logits[np.arange(len(labels)), labels] -= 1
    grad_logits /= len(labels)

    grad_W2 = cache["H"].T @ grad_logits
    grad_b2 = grad_logits.sum(axis=0)
    grad_hidden = grad_logits @ params["W2"].T
    grad_z1 = grad_hidden * (cache["Z1"] > 0)
    return {
        "W1": cache["X"].T @ grad_z1,
        "b1": grad_z1.sum(axis=0),
        "W2": grad_W2,
        "b2": grad_b2,
    }

14.3 Initialization and training

If $z_j = \sum_{i=1}^{n} w_{ij}x_i$ and the factors are independent with zero means, then

\Var[z_j] = \sum_i \Var[w_{ij}x_i] = n\,\Var[w] \Var[x].\tag{14.5}

Keeping activation variance stable therefore suggests $\Var[w] \approx 1/n$ . Xavier initialization balances the forward and backward directions for symmetric activations [glorot2010]. ReLU drops roughly half the signal, so He initialization doubles the variance to $2/n$ for ReLU layers [he2015delving]; that is why the first layer above uses $\sqrt{2/5}$ .

This derivation is deliberately approximate. Real minibatches are not independent, learned weights quickly stop being independent of activations, and nonlinearities change more than a single variance. The point is still practical: start with a scale that neither crushes signals to zero nor blows them up before the first update. A bad scale can make a correct gradient formula look broken because all rows predict the same class or all ReLUs are inactive.

Training repeats forward pass, backward pass, parameter update. The notebook uses Adam so it can learn quickly in the browser; Adam itself is derived in Chapter 15. The listing below intentionally uses full-batch SGD so the mechanism is visible: subtract a small multiple of every gradient from its matching parameter.

Listing 14.3 A tiny plain-SGD training run

def synthetic_data(seed=1, examples_per_class=24):
    rng = np.random.default_rng(seed)
    centers = np.array([[-1.5, -1.5, 1, 0, -1], [-1.5, 1.5, -1, 1, 0],
                        [1.5, -1.5, 0, -1, 1], [1.5, 1.5, 1, 1, 1]],
                       dtype=np.float32)
    labels = np.repeat(np.arange(4), examples_per_class)
    points = np.concatenate([
        center + rng.normal(0, 0.55, size=(examples_per_class, 5))
        for center in centers
    ]).astype(np.float32)
    return points, labels


def train_sgd(inputs, labels, epochs=80, learning_rate=0.08, seed=0):
    params = initialize(seed)
    history = []
    for _ in range(epochs):
        loss, cache = forward(params, inputs, labels)
        history.append(float(loss))
        gradients = backward(params, cache, labels)
        for name, value in params.items():
            value -= learning_rate * gradients[name]
    return params, np.array(history)

A network with one hidden nonlinear layer can approximate any continuous function on a compact set, given enough hidden units. That theorem is about existence; it does not promise that SGD will find the approximation or that the learned rule will generalize.

14.4 Overfitting

The companion notebook makes overfitting concrete. It first trains the 5 → 10 → 4 network on clean synthetic labels. Then it keeps only 16 examples, randomly reassigns balanced labels, resets the weights and Adam state, and trains much longer. The network can memorize the tiny noisy set, but validation accuracy stays poor and validation loss can rise as wrong predictions become confident.

That experiment separates optimization from generalization. A low training loss only says the parameters fit the examples used for updates. Validation data estimates whether the learned rule transfers to unseen examples; regularization, data scale, early stopping, and model architecture all change that gap [goodfellow2016].

The model is small enough that memorization looks surprising, but the parameter count is not the whole story. ReLU networks carve the input space into many linear regions, and a tiny data set leaves most of those regions unconstrained. Larger validation sets, fresh test sets, and ablations are how we detect that a training curve has become a memory of examples rather than a useful classifier.

In practice

Modern language models are enormous descendants of this computation: affine projections, nonlinear feed-forward blocks, softmax losses, and backpropagation through the whole stack. Transformers add attention and residual paths [vaswani2017attention], and recent LLMs often use gated feed-forward activations rather than ReLU [grattafiori2024]. The 5 → 10 → 4 MLP is still the right microscope because the matrix gradients and initialization logic are the same.

Key equations

\mZ_1 = \mX\mW_1 + \vb_1,\quad \mH = \max(0, \mZ_1)

L = -\frac{1}{B}\sum_i \log P_{i,y_i}

\bar{\mZ}_2 = (\mP - \mY)/B

\bar{\mW} = \mA^\T\bar{\mZ},\quad \bar{\mA}=\bar{\mZ}\mW^\T

\Var\!\left[\sum_i w_i x_i\right] = n\,\Var[w]\Var[x]

14.5 Teach it

The one-sentence version: an MLP alternates matrix multiplies with nonlinearities, then uses backpropagation to assign each parameter its share of the loss. Analogy: each layer reshapes space, and the gradient tells every handle which direction would make the final label score better. Board steps: (1) write $\mX\mW_1 + \vb_1$ , ReLU, $\mH\mW_2 + \vb_2$ ; (2) turn logits into probabilities and loss; (3) start backward with $(\mP-\mY)/B$ ; (4) use $\mA^\T\bar{\mZ}$ and the ReLU mask. Misconceptions: softmax is not needed during prediction if only the argmax matters, but it is needed for probabilities and the loss; validation loss is not optimized directly; universal approximation does not prevent overfitting. Check for understanding: if the hidden width changes from 10 to 20, which parameter shapes and gradient shapes change?

14.6 Exercises

Exercise 14.1 ★ Shapes and parameter count

For a batch $\mX \in \R^{B \times 5}$ , list the shapes of $\mZ_1$ , $\mH$ , $\mZ_2$ , and each parameter in the 5 → 10 → 4 network. How many trainable scalars are there?

Solution in Appendix D

Exercise 14.2 ★★ Softmax-cross-entropy gradient

Starting from $p_k = e^{z_k}/\sum_j e^{z_j}$ and $\ell = -\sum_k y_k \log p_k$ , derive $\partial \ell/\partial z_j = p_j - y_j$ . Why does the batch implementation divide by $B$ exactly once?

Solution in Appendix D

Exercise 14.3 ★★ Variance propagation

Assume $z = \sum_{i=1}^{n} w_i x_i$ , with independent zero-mean factors and common variances $\Var[w]$ and $\Var[x]$ . Derive (14.5) and explain the ReLU change that leads to He initialization.

Solution in Appendix D

Exercise 14.4 ★★★ Implement and check

Implement the listing functions for a tiny MLP, verify the gradients with finite differences, and train the synthetic data with full-batch SGD. Then repeat on a tiny noisy subset and report why training accuracy and clean-label accuracy can separate.

Solution in Appendix D

References

[goodfellow2016] I. Goodfellow, Y. Bengio, and A. Courville. Deep Learning. MIT Press, 2016. https://www.deeplearningbook.org
[grattafiori2024] A. Grattafiori et al. The Llama 3 herd of models. 2024. arXiv:2407.21783
[rumelhart1986] D. E. Rumelhart, G. E. Hinton, and R. J. Williams. Learning representations by back-propagating errors. Nature 323, 533–536, 1986.
[glorot2010] X. Glorot and Y. Bengio. Understanding the difficulty of training deep feedforward neural networks. AISTATS 2010.
[he2015delving] K. He et al. Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification. 2015. arXiv:1502.01852
[vaswani2017attention] A. Vaswani et al. Attention Is All You Need. 2017. arXiv:1706.03762