# Module 4: How a Model Learns From Data

> Train a small neural network end to end: neurons and MLPs, softmax with cross-entropy, overfitting, hyperparameters, and how to read an experiment.

- Part I: Foundations · core track · Foundational · 5 hours
- By the end you can train a tiny MLP end to end, diagnose overfitting, and read an experiment's result critically.
- Builds on: [Module 3: Probability and Information](https://transformers-and-math.drmhse.com/md/03-probability-and-information.md)
- Web page (interactive figures, labs and equation explainers): https://transformers-and-math.drmhse.com/modules/04-how-a-model-learns
- From *Transformers & Math* by Mike Chumba: https://transformers-and-math.drmhse.com

Modules 2 and 3 supplied the two halves of learning: a loss that measures how wrong a model is, and gradient descent to make it less wrong. This module puts them together on a small neural network, then adds what makes training a science: why fitting the training data is not the goal, which knobs you choose by hand, and how to tell a real improvement from luck.

## 4.1 Supervised and Self-Supervised Learning

Most machine learning is *supervised*: each example pairs an input with the wanted answer, its *label* (an image labelled `cat`, an email labelled `spam`), and the loss compares the model's output with the label.

Labels are expensive, so labelled datasets hold millions of examples, not trillions. Language models make labels from the data itself: hide the next token and predict it from the tokens before. A 20-token sentence is 20 training examples, and the whole internet becomes a labelled dataset. This is *self-supervised learning*, and it is why language models could grow so large: the training signal grows with the supply of text. Section 8.4 writes the objective out, and section 8.2 shows a second self-supervised task, filling in blanks.

A third kind, *reinforcement learning*, has no labels, only a score for how good an output was. Module 21 uses it to tune assistants to people's preferences.

## 4.2 Neurons, Layers, and the MLP

A *neuron* is the smallest unit of a neural network. It takes a vector of inputs, computes a weighted sum plus a constant, and passes the result through a fixed nonlinear function:

$$
a = \phi(\mathbf{x}\mathbf{w} + b)
$$

The weights $\mathbf{w}$ and *bias* $b$ are learned. $\phi$ is the *activation function*: ReLU, $\max(0, z)$, zeroes negative values; $\tanh$ squashes into $(-1, 1)$; GELU is a smooth ReLU used in transformers (section 7.4).

A *layer* is many neurons reading the same input. With their weight vectors as the columns of $W$, a layer is one matrix multiplication followed by $\phi$ entry by entry: $\mathbf{a} = \phi(\mathbf{x}W + \mathbf{b})$. A *multi-layer perceptron* (MLP) chains such layers. The layers between input and output are *hidden layers*, and their size is the network's *width*.

### Why the Nonlinearity Matters
Without $\phi$, two layers give $(\mathbf{x}W_1 + \mathbf{b}_1)W_2 + \mathbf{b}_2 = \mathbf{x}(W_1W_2) + (\mathbf{b}_1 W_2 + \mathbf{b}_2)$: one linear layer again. Any stack of linear layers can only draw straight boundaries; the nonlinearity is what lets depth add power. With enough hidden neurons, one hidden layer can approximate any reasonable function (the *universal approximation theorem*), but depth makes many functions far cheaper to represent.

Every transformer block contains an MLP, the *feed-forward network* (section 7.4), which holds most of the model's parameters.

## 4.3 Logits, Softmax, and Cross-Entropy

A model produces raw scores called *logits*, $\mathbf{z} = [z_1, \dots, z_N]$, any real numbers. Softmax turns them into probabilities: positive, and summing to 1.

$$
\operatorname{Softmax}(\mathbf{z})_i = \frac{e^{z_i}}{\sum_{j=1}^N e^{z_j}}
$$

### Numerical Catastrophe: Floating-Point Overflow
Floating-point formats have a largest value; anything bigger becomes infinity (`inf`). Because $e^{z}$ grows so fast, ordinary logits can overflow:

- In 32-bit floats (`fp32`) the largest value is about $3.4 \times 10^{38} \approx e^{88.72}$, so $e^{z}$ overflows once $z > 88.72$.
- In 16-bit half precision (`fp16`) the largest value is $65{,}504 \approx e^{11.09}$, so $e^{z}$ overflows once $z > 11.09$.
- `bf16` keeps fp32's exponent range and gives up precision instead, so it overflows near the same point as fp32.

Once any $e^{z_i}$ is `inf`, softmax computes $\infty / \infty$ = `NaN` ("not a number"). Backpropagation carries it into the gradients and then the weights, and any arithmetic with `NaN` gives `NaN`.

### Shift-Invariance & Safe Softmax
Subtracting the same number $c$ from every logit leaves the softmax unchanged:

$$
\operatorname{Softmax}(\mathbf{z} - c)_i = \frac{e^{z_i - c}}{\sum_j e^{z_j - c}} = \frac{e^{-c} e^{z_i}}{e^{-c} \sum_j e^{z_j}} = \operatorname{Softmax}(\mathbf{z})_i
$$

The factor $e^{-c}$ cancels top and bottom. Standard implementations (PyTorch, FlashAttention, TensorRT-LLM) use $c = \max_j (z_j)$:

$$
\operatorname{SafeSoftmax}(\mathbf{z})_i = \frac{e^{z_i - \max_k(z_k)}}{\sum_{j=1}^N e^{z_j - \max_k(z_k)}}
$$

Every exponent is now zero or negative, so every exponential lies in $(0, 1]$ and nothing can overflow.

```mermaid
graph LR
    Z["Raw logits z_i (can be as large as 120)"] --> Max["Find m = max(z_k)"]
    Max --> Sub["Subtract: z_i − m ≤ 0"]
    Sub --> Exp["exp(z_i − m) lies in (0, 1]: cannot overflow"]
    Exp --> Sum["Divide each by the sum"]
    Sum --> P["Probabilities p_i, identical to plain softmax"]
```

> **In words.** Subtract the largest logit from every logit, then exponentiate and divide by the sum. The answer is unchanged; only overflow is prevented.

### Temperature Scaling Regimes
When sampling text, the logits are often divided by a *temperature* $T > 0$ first:

$$
p_i = \frac{\exp(z_i / T)}{\sum_j \exp(z_j / T)}
$$

- **$T = 1$:** the model's own probabilities.
- **$T \to 0$ (greedy decoding, always picking the most likely token):** every gap $\frac{z_{\max} - z_i}{T}$ grows without bound, so all probability lands on the largest logit, a *one-hot* distribution. $\arg\max_k z_k$ is the position of the largest logit:
  $$
  \lim_{T \to 0^+} p_i = \begin{cases} 1, & \text{if } i = \arg\max_k z_k \\ 0, & \text{otherwise} \end{cases}
  $$
- **$T \to \infty$ (uniform):** every $z_i / T$ goes to $0$, so every token gets the same weight:
  $$
  \lim_{T \to \infty} p_i = \frac{1}{N}
  $$

> **Interactive figure** `softmax-temperature`, on the web page: https://transformers-and-math.drmhse.com/modules/04-how-a-model-learns#figure-softmax-temperature
>
> **Try:** Before moving the slider, decide where the point must end up as the temperature grows very large, and why the answer does not depend on the logits at all.
>
> **What happens:** At the centre, where every probability is $\tfrac13$. Dividing by a huge $T$ shrinks every logit toward 0, so every exponential approaches $e^0 = 1$ and the differences between the logits stop mattering.

> **In words.** With three logits, every softmax output is a point in a triangle whose corners are the one-hot distributions. Lowering the temperature slides the point toward the largest logit's corner (greedy decoding); raising it slides it to the center, where each probability is a third.

### Softmax Jacobian and Gradient Dynamics
How does each probability $p_i = \operatorname{Softmax}(\mathbf{z})_i$ change when logit $z_j$ changes?

$$
\frac{\partial p_i}{\partial z_j} = p_i (\delta_{ij} - p_j) = \begin{cases} p_i(1 - p_i), & \text{if } i = j \\ -p_i p_j, & \text{if } i \neq j \end{cases}
$$

$\delta_{ij}$ is $1$ when $i = j$ and $0$ otherwise. This Jacobian is not diagonal: raising one logit lowers every other probability, since they must still sum to 1. Both cases follow from the quotient rule with $S = \sum_k e^{z_k}$: for $i = j$, $\frac{e^{z_i} S - e^{z_i} e^{z_i}}{S^2} = p_i - p_i^2$; for $i \neq j$, $\frac{-e^{z_i} e^{z_j}}{S^2} = -p_i p_j$.

If one logit dominates, $p_i \approx 1$ and every other $p_j \approx 0$, so every entry of the Jacobian is nearly zero and almost no gradient flows back. The softmax is *saturated*, and learning through it stalls. This is one reason attention divides its scores by $\sqrt{d_k}$.

> **Measured** (lab 1b, `experiments/softmax_overflow.py`). The size of the Jacobian (its Frobenius norm) drops from $0.433$ for equal logits to $0.000193$ for saturated ones, a factor of $2{,}240$, as the largest probability goes from $0.25$ to $0.9999$. Only the two ends are compared: along some paths between them the norm rises before it falls, which is why section 6.3 tracks the top class's $p(1-p)$ instead.

> **Measured** (lab 1b). exp() overflows fp32 at $z > 88.72$ and fp16 at $z > 11.09$, matching the values quoted in the overflow section. Passing $z = [120, 0, -120]$ to the plain softmax formula returns `nan`, while the safe version gives the exact answer. Note which code is at risk: `torch.softmax` already subtracts the row maximum internally, so the model code shipped with this course never overflows. The danger is in writing the formula out by hand, not in calling the library.

### Softmax With Cross-Entropy: The Gradient That Trains Every LLM
Saturation matters when the softmax output feeds further computation, as in attention. At a model's output, the softmax feeds the cross-entropy loss $\mathcal{L} = -\log p_c$ for the correct class $c$ (section 3.5), and differentiating through both at once gives something much simpler:

$$
\frac{\partial \mathcal{L}}{\partial z_j} = p_j - y_j
$$

$y_j$ is $1$ for the correct class and $0$ otherwise. In one line: $-\log p_c = -z_c + \log \sum_k e^{z_k}$, and the second term's derivative is $p_j$.

The gradient is "what the model predicted, minus what was right". When the model is confidently wrong ($p_c \approx 0$), the correct class gets a gradient near $-1$, the strongest push; it vanishes only when the model is confidently right. The log cancels the exponential, which is why a saturated output layer is no problem while a saturated softmax inside attention is.

## 4.4 Generalization, Overfitting, and Regularization

The goal is not a low loss on the training data but on data the model has never seen: *generalization*. A model that memorized its examples would be useless on anything else.

### Three Splits
To measure generalization honestly, the data is divided before training:

- The **training set** is what gradient descent sees.
- The **validation set** is held out and used for decisions: learning rate, model size, when to stop. Its loss is the *validation loss*.
- The **test set** is touched once, at the end, to report the final number.

Making decisions on the test set turns it into a second validation set, and the reported number becomes optimistic. Section 13.3 shows the same failure with public benchmarks.

### Overfitting and Underfitting
Watch the training and validation losses together:

- **Underfitting.** Both are high, still falling or stalled. The model is too small or undertrained.
- **A good fit.** Both fall, and stay close.
- **Overfitting.** Training loss keeps falling while validation loss turns and rises: the model is fitting noise in its particular examples.

Overfitting needs a model big enough to memorize and data small enough to be memorized. Large language models see almost every example once and rarely overfit; small models trained for many passes over small datasets overfit readily.

### Regularization
A *regularizer* discourages memorization at a small cost to the training loss:

- **More data**, the most effective of all.
- **Weight decay** pulls every weight slightly toward zero each step, discouraging the large, finely tuned weights memorization needs (section 10.2).
- **Dropout** zeroes a random fraction of activations at each training step, so the model cannot rely on any single feature (Srivastava et al., 2014). GPT-2 used $0.1$; most modern pretraining uses zero, since it rarely overfits.
- **Early stopping** keeps the weights from the step with the lowest validation loss.

## 4.5 Hyperparameters, Batches, Steps, and Epochs

*Parameters* are what gradient descent learns. Everything chosen by hand is a *hyperparameter*: learning rate, number and width of layers, batch size, weight decay, dropout, training length. They are chosen by comparing validation loss across several values.

- A **batch** is the set of examples for one gradient computation. Larger batches give a less noisy gradient (section 3.2) at a higher cost per step.
- A **step** is one weight update: loss on one batch, backpropagate, gradient descent.
- An **epoch** is one full pass over the training set: $n / B$ steps for $n$ examples and batch size $B$.

Small models train for many epochs. Large language models train for about one epoch or less, so progress is counted in *tokens seen* instead (section 10.5).

### Reading a Loss Curve
A *loss curve* plots loss against step. A healthy one falls fast, then flattens. A flat curve suggests a learning rate too small, or a bug disconnecting loss from weights. Sudden spikes suggest a learning rate too large or overflow (section 4.3); section 10.4 lists the defences.

## 4.6 Reading an Experiment

Rerun any experiment and its numbers come out slightly different. Reading results well means knowing why.

### Seeds and Run-to-Run Noise
Initial weights, batch order and dropout are drawn from a *pseudo-random* generator, which repeats exactly from the same *seed*. A different seed gives a different but equally valid run, and the final loss moves a little: *run-to-run noise*.

A difference between methods means something only if it is clearly larger than the noise. If changing the seed moves validation loss by $0.02$, a single-run improvement of $0.015$ shows nothing yet. Run each method with several seeds and compare the spread.

### Ablations
An *ablation* changes one component and keeps everything else (data, seed, steps) identical, isolating that component's effect. Most comparisons in this course are ablations: sinusoidal against learned positions (section 7.1), Post-LN against Pre-LN (section 8.5).

### Reading a Measured Box
When you read a **Measured** box, ask four questions:

1. **What was measured?** A derivation checked to floating-point precision (such as an identity holding to $10^{-7}$) is a different kind of claim from a training comparison.
2. **At what scale?** A result on a model with 100,000 parameters may not hold at a billion. The boxes say when this matters.
3. **How many seeds?** One seed shows that something *can* happen, not how reliably.
4. **Compared with what noise?** A difference smaller than the run-to-run noise is not a result.

The **Caveat** boxes exist to answer these questions in advance.

## 4.7 Worked Example: Training a Tiny MLP

The figure trains an MLP in your browser to separate two classes of points in a plane: two coordinates in, two logits out, softmax and cross-entropy, then backpropagation and gradient descent on each batch.

> **Interactive figure** `mlp-playground`, on the web page: https://transformers-and-math.drmhse.com/modules/04-how-a-model-learns#figure-mlp-playground
>
> **Try:** Set the training set to its smallest size and the hidden layer to its largest width, then train. Before you press play, predict what the two loss curves will do.
>
> **What happens:** The training loss falls nearly to zero while the validation loss falls, bottoms out and then climbs: overfitting. The decision boundary wraps itself around individual training points. Add data, shrink the network, or add weight decay, and the two curves stay together.

> **In words.** A plane of two-coloured points, shaded by the network's current prediction, beside a chart of training and validation loss. Controls set the number of points, network width and depth, activation, learning rate and weight decay. As training runs, the shading reshapes around the data and the curves show underfitting, fitting or overfitting.

Try, in order: with ReLU and moderate width, the boundary bends to follow the data. Switch the activation to *none* and it stays a straight line however many layers you add (section 4.2). Raise the learning rate and the loss curve jumps around (section 2.4). Finally, make it overfit with the smallest training set and widest network, then cure it.

> **Browser lab** `mlp-overfit`: an experiment that runs in the reader's browser on the web page: https://transformers-and-math.drmhse.com/modules/04-how-a-model-learns#lab-mlp-overfit

## Exercises

**Exercise (check).** Why is next-token prediction called *self-supervised* rather than supervised?

*Solution:* Because the label for every example, the next token, comes from the text itself. Nobody has to annotate anything, so every piece of text supplies as many training examples as it has tokens.

**Exercise (apply).** Show that two linear layers with no activation between them, $\mathbf{x}W_1W_2$ with $W_1 \in \mathbb{R}^{4 \times 16}$ and $W_2 \in \mathbb{R}^{16 \times 3}$, are equivalent to one linear layer. What is that layer's shape, and what is the largest rank it can have?

*Answer:* 3 (within 0)

*Solution:* $\mathbf{x}W_1W_2 = \mathbf{x}(W_1W_2)$ and $W_1W_2$ is a single $4 \times 3$ matrix. Its rank is at most $\min(4, 3) = 3$. The 16-wide hidden layer adds nothing without a nonlinearity.

**Exercise (apply).** A model's logits are $[2, 1, 0]$ and the correct class is the first. Compute the gradient of the cross-entropy loss with respect to the first logit.

*Answer:* -0.3348 (within 0.001)

*Solution:* Softmax gives $p_1 = e^2 / (e^2 + e^1 + e^0) = 7.389 / 11.107 \approx 0.665$. The gradient is $p_1 - y_1 = 0.665 - 1 \approx -0.335$: push the correct logit up.

**Exercise (apply).** A training set has 60,000 examples and the batch size is 128. How many steps is 3 epochs, rounding each epoch up to a whole number of steps?

*Answer:* 1407 (within 0)

*Solution:* One epoch is $\lceil 60{,}000 / 128 \rceil = 469$ steps, so three epochs are $1{,}407$ steps.

**Exercise (apply).** Describe what the training and validation loss curves look like when a model overfits, and name two ways to reduce it.

*Solution:* The training loss keeps falling while the validation loss stops falling and starts to rise. Remedies include more training data, a smaller model, weight decay, dropout, and early stopping at the step with the lowest validation loss.

**Exercise (extend).** Two methods are compared with one seed each. Method B's validation loss is $0.015$ lower than method A's. Three more seeds of method A give losses spread over a range of $0.04$. What can you conclude, and what would you do next?

*Solution:* Nothing yet: the difference is smaller than method A's own run-to-run spread, so it could easily be noise. Run both methods over several seeds and compare the distributions (for example their means and standard deviations); only a gap clearly larger than the spread supports a claim.

**Exercise (extend).** Why does a saturated softmax stall learning inside attention but not at the output layer of a classifier trained with cross-entropy?

*Solution:* At the output, the loss is $-\log p_c$, and the log cancels the exponential: the gradient with respect to the logits is $\mathbf{p} - \mathbf{y}$, which is large whenever the model is confidently wrong. Inside attention, the softmax output feeds further computation with no log to cancel it, so the gradient passes through the softmax Jacobian, whose entries $p_i(\delta_{ij} - p_j)$ are all near zero when the softmax is saturated.

### Core Academic References
- **Goodfellow, Bengio & Courville (2016)**: *Deep Learning*, chapters 5–7 (Machine Learning Basics, Deep Feedforward Networks, Regularization). MIT Press.
- **Srivastava et al. (2014)**: *Dropout: A Simple Way to Prevent Neural Networks from Overfitting*. JMLR 15: 1929–1958.
- **Hornik, Stinchcombe & White (1989)**: *Multilayer Feedforward Networks are Universal Approximators*. Neural Networks, 2(5): 359–366.
- **Bouthillier et al. (2021)**: *Accounting for Variance in Machine Learning Benchmarks*. MLSys 2021. [arXiv:2103.03098](https://arxiv.org/abs/2103.03098)