How a Model Learns From Data

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 MLPRun this section

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:

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-EntropyRun this section

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.

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:

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

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

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:

• $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}$$

  What does this say?introduces 2 notations
• $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}$$

  What does this say?

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?

$\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:

$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 RegularizationRun this section

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 ExperimentRun this section

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 MLPRun this section

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.

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.

Exercises

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