Training and Evaluation

The Training Problem in One Equation

In section 2 we built a bidirectional LSTM that maps a sentence to two logits. Right now its parameters are random; it is wrong on roughly half of the validation reviews. Training is the process that turns that untrained architecture into a useful sentiment classifier.

Every supervised learning problem — for an LSTM, a transformer, or a logistic regression — reduces to exactly one optimization:

$\theta^{\star} = \arg\min_{\theta} \; \mathcal{L}(\theta) = \arg\min_{\theta} \; \frac{1}{N} \sum_{i=1}^{N} \ell\big(f_{\theta}(x_i),\; y_i\big)$

Read this equation word by word. $\theta$ is every learnable number in the network (embeddings, LSTM gate weights, linear head, biases — about 1.8 million of them for our classifier). $f_{\theta}$ is the model as a function of those parameters. $\ell$ is a per-example loss — how wrong the model is on a single review. The full loss $\mathcal{L}$ averages that across all $N$ training examples, and we want the $\theta$ that makes it as small as possible.

The central claim of deep learning. If you pick a flexible enough $f_{\theta}$, a good enough $\ell$, and enough data, then repeatedly nudging $\theta$ opposite the gradient of $\mathcal{L}$ is sufficient to discover parameters that generalize to unseen inputs. No explicit feature engineering, no rulebook.

Nothing in this equation cares whether $f_{\theta}$ is a two-parameter logistic regression or a 70-billion-parameter transformer. The machinery — dataset, forward pass, loss, backward pass, optimizer step — is identical. Scale is what changes, and we'll trace every one of those scaling pressures down to Flash Attention and the KV-cache by the end of this section.

The Dataset and Mini-Batches

Our dataset is the tokenised, padded reviews we prepared in Section 18.1 — say 25 000 labelled movie reviews split into 20 000 for training, 2 500 for validation, and 2 500 held out for the final test. In PyTorch this is two Dataset objects wrapped by DataLoaders that yield shuffled mini-batches.

Why mini-batches instead of feeding all 20 000 reviews at once? Three reasons, in decreasing order of how often they matter:

1. Memory. A single batch of 64 reviews at length 200 with a BiLSTM is already ~100 MB of activations. The full 20 000-review gradient would need hundreds of gigabytes. A GPU cannot hold it.
2. Noise is useful. A mini-batch gradient is a noisy estimate of the full-dataset gradient. The noise helps the optimizer escape narrow, sharp minima that do not generalize.
3. Parallelism. 64 reviews fit perfectly on a modern GPU's SIMD units. Larger batches stop scaling when the layer is already at peak throughput.

The animation below shows one epoch through the DataLoader. Each step pulls a batch of 64 from the shuffled pool; once the pool is empty, the epoch ends, the pool refills, and a fresh permutation is drawn. Turning shuffle off reveals why the default is on: consecutive batches become correlated and the gradient updates lose their SGD character.

An epoch is one full pass of the DataLoader over the training set — 20 000 / 64 ≈ 313 steps of SGD per epoch. A typical sentiment run completes in 20-30 epochs, or roughly 7 000 parameter updates.

The Loss Function: Cross-Entropy

For a classification problem the standard per-example loss is cross-entropy. For a single sample with true class $y \in \{1, \dots, C\}$ and the model's predicted distribution $p_{\theta}(c \mid x)$ it is:

$\ell_{\text{CE}}(f_{\theta}(x), y) = -\log p_{\theta}(y \mid x) = -\log \frac{e^{z_y}}{\sum_{c=1}^{C} e^{z_c}}$

where $z_c$ is the raw logit the model produced for class $c$. For the binary sentiment case $C = 2$ and we can equivalently write $\ell_{\text{BCE}} = -[y \log p + (1-y) \log(1-p)]$ where $p = \sigma(z_1 - z_0)$.

Why this particular loss?

Cross-entropy arises naturally from three different angles that all land on the same formula:

1. Maximum likelihood. Under the model's own distribution, the probability of observing the training labels is $\prod_i p_{\theta}(y_i \mid x_i)$. Taking the negative log gives precisely the cross-entropy sum. Minimising cross-entropy is maximising the likelihood of the data.
2. Information theory. Cross-entropy measures the extra bits needed to encode a sample from the true distribution using the model's distribution. When the two match, $H(p^{\star}, p_{\theta}) = H(p^{\star})$ (the irreducible entropy of the data) and the excess term drops to zero.
3. Well-behaved gradients. Combined with softmax, its gradient collapses to the clean $\partial \ell / \partial z_c = p_c - \mathbb{1}[c = y]$. No saturating derivatives, no vanishing gradients at wrong-but- confident predictions (MSE does vanish there).

This gradient simplification is not a happy accident — it is the reason every classifier in this book, from logistic regression to transformer language models, uses the softmax/cross-entropy pair. We will exploit it directly in the NumPy trace below.

From SGD to Adam: Learning Rates That Adapt

Once we have $\mathcal{L}(\theta)$ and its gradient, the update rule is deceptively simple:

$\theta_{t+1} = \theta_t - \eta \; \nabla_{\theta} \mathcal{L}(\theta_t)$

This is stochastic gradient descent (SGD). A single scalar learning rate $\eta$ controls step size for every parameter. It works, but on real networks it is fragile: one learning rate has to simultaneously suit parameters whose gradients vary by many orders of magnitude. The embedding matrix, the LSTM gate weights and the final bias all live on different scales.

Momentum — remember the past

First fix: accumulate an exponential moving average of the gradients themselves.

$m_t = \beta_1 \, m_{t-1} + (1 - \beta_1) \, g_t, \qquad \theta_{t+1} = \theta_t - \eta \, m_t$

Now consistent gradient directions build up speed, while jittery coordinates (where gradients flip sign each step) average toward zero. This is the ball-rolling-down-a-hill intuition.

Per-parameter scaling — Adam

Adam (Kingma & Ba, 2014) adds a second EMA, this time of squared gradients, and uses it to scale each parameter's step individually:

$m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t, \quad v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t^2$

$\hat{m}_t = m_t / (1 - \beta_1^t), \quad \hat{v}_t = v_t / (1 - \beta_2^t)$

$\theta_{t+1} = \theta_t - \eta \, \hat{m}_t / (\sqrt{\hat{v}_t} + \varepsilon)$

Where do these weird terms come from?

• $m_t$ — momentum, the mean direction of recent gradients. Typical $\beta_1 = 0.9$ gives an effective memory of about 10 steps.
• $v_t$ — the mean-squared gradient per parameter. Large $v$ means the parameter has been receiving big gradients; Adam shrinks its step. Small $v$ means the parameter has been receiving tiny gradients; Adam grows its step.
• $\hat{m}, \hat{v}$ — bias-corrected versions. Since $m_0 = v_0 = 0$, the raw EMAs are biased toward zero in early steps; dividing by $1 - \beta^{t}$ undoes this.
• $\varepsilon \approx 10^{-8}$ — prevents division by zero when a parameter has received no gradient yet.

The effective step size for each parameter is $\eta / \sqrt{\hat{v}_t}$: small-gradient parameters get bigger steps, large-gradient parameters get smaller steps. Adam is the universal default optimizer of modern deep learning precisely because this per-parameter normalization makes the single global learning rate robust.

Interactive: SGD vs Adam on a Narrow Valley

To see why Adam's per-parameter scaling matters, here is a 3D loss landscape with two optimizers dropped at the same starting point. The loss is $L(w_1, w_2) = \tfrac{1}{2}(w_1^2 + 9 w_2^2)$ — a textbook narrow valley: nine times steeper along $w_2$ than along $w_1$. Real loss surfaces have many such direction-dependent curvatures; this one simply isolates the phenomenon.

Click Play. The red ball (SGD, fixed learning rate) is caught immediately by the steep $w_2$ direction: its step is too large for that curvature, so it overshoots, lands on the opposite wall, overshoots again, and zig-zags its way toward the minimum. Meanwhile it makes only slow progress along the gentler $w_1$ direction.

The blue ball (Adam) tracks $v_t$ separately for $w_1$ and $w_2$. On$w_2$ the running squared gradient is large, so the effective step $\eta / \sqrt{\hat{v}}$ shrinks — no overshoot. On $w_1$ it stays small, so Adam pushes harder. The result is a smooth, almost straight trajectory to the minimum.

The Five-Step Training Loop

Every training step — every one of the ~7 000 weight updates our LSTM will perform — is exactly the same five operations. Learn this rhythm once and it applies to every model you will ever train.

1. Forward pass. $\hat{y} = f_{\theta}(x)$ — feed the batch through the model and record every intermediate activation.
2. Loss. $\ell = \mathcal{L}(\hat{y}, y)$ — a single scalar tensor measuring current wrongness.
3. Backward pass. $\nabla_{\theta} \ell$ — the chain rule, reverse-mode, filling every parameter's .grad attribute.
4. Optimizer step. $\theta \leftarrow \mathrm{Update}(\theta, \nabla_{\theta} \ell)$ — Adam applies its moment-corrected rule.
5. Zero the gradients. optimizer.zero_grad() — clear the accumulators so the next .backward() starts from a clean slate.

Periodically — usually once per epoch — we pause training and run the full validation set through the forward pass only, recording loss and accuracy. That gives us the two curves in the visualization below.

Tracing LSTM Backprop in NumPy

The logistic-regression trace in the next section is the cleanest way to see a gradient derivation on one page, but you also need to know that the LSTM step from section 18.2 has a concrete backward pass. Below is a single cell traced line-by-line — forward, then backward — so you can see that $\partial L / \partial c_t$ passes straight through the additive identity (the reason gradients do not vanish) while each gate's weight gradient comes from its own sigmoid or tanh derivative.

         d_embed=0   d_embed=1
[i,0]        0.5         0.3
[i,1]        0.1        -0.2
[f,0]        0.4         0.6
[f,1]        0.2         0.1
[g,0]        0.7         0.5
[g,1]        0.3         0.4
[o,0]        0.2        -0.1
[o,1]        0.0         0.3

All zeros.

           d_embed=0   d_embed=1
[i,0]        0.0018      0.0014
[i,1]       -0.0024     -0.0018
[f,0]        0.0009      0.0007
[f,1]       -0.0010     -0.0008
[g,0]        0.0087      0.0065
[g,1]       -0.0190     -0.0142
[o,0]        0.0033      0.0025
[o,1]       -0.0033     -0.0025

1import numpy as np
2
3# --- From section 18.2: a single LSTM step, plus its gradient ---
4d_embed, d_hidden = 2, 2
5
6# Inputs at time t (one step only, to keep values on one screen)
7x_t     = np.array([0.40, 0.30])           # embedding for "movie"
8h_prev  = np.array([0.10, -0.05])
9c_prev  = np.array([0.20,  0.10])
10
11# Gate pre-activation matrices (reused from section 18.2)
12W_x = np.array([
13    [ 0.5,  0.3], [ 0.1, -0.2],            # input gate
14    [ 0.4,  0.6], [ 0.2,  0.1],            # forget gate
15    [ 0.7,  0.5], [ 0.3,  0.4],            # cell candidate
16    [ 0.2, -0.1], [ 0.0,  0.3],            # output gate
17])
18W_h = np.zeros((8, 2))
19b   = np.zeros(8)
20
21def sigmoid(z): return 1.0 / (1.0 + np.exp(-z))
22
23# --- Forward ---
24a      = W_x @ x_t + W_h @ h_prev + b
25i_g    = sigmoid(a[0:2])
26f_g    = sigmoid(a[2:4])
27g_c    = np.tanh(a[4:6])
28o_g    = sigmoid(a[6:8])
29c_t    = f_g * c_prev + i_g * g_c
30tanh_c = np.tanh(c_t)
31h_t    = o_g * tanh_c
32
33# --- Upstream gradient: pretend dL/dh_t comes from the next layer ---
34dL_dh = np.array([0.10, -0.20])
35
36# --- Manual backward through the LSTM cell ---
37do_g   = dL_dh * tanh_c                    # dL/do = dL/dh * tanh(c)
38dc_t   = dL_dh * o_g * (1 - tanh_c**2)     # dL/dc = dL/dh * o * (1-tanh^2(c))
39df_g   = dc_t  * c_prev                    # dL/df
40di_g   = dc_t  * g_c                       # dL/di
41dg_c   = dc_t  * i_g                       # dL/dg
42
43# Through each gate's activation (sigmoid/tanh derivatives)
44da_i   = di_g * i_g * (1 - i_g)
45da_f   = df_g * f_g * (1 - f_g)
46da_g   = dg_c * (1 - g_c**2)
47da_o   = do_g * o_g * (1 - o_g)
48da     = np.concatenate([da_i, da_f, da_g, da_o])     # (8,)
49
50# Gradient to W_x: outer product of da and x_t
51dW_x   = np.outer(da, x_t)                 # shape (8, 2)
52db     = da                                # shape (8,)
53dh_prev = W_h.T @ da                       # shape (2,)
54
55print("h_t       :", np.round(h_t,   4))
56print("c_t       :", np.round(c_t,   4))
57print("dW_x[0,:] :", np.round(dW_x[0], 4))  # input-gate row 0
58print("db        :", np.round(db,    4))
59print("dh_prev   :", np.round(dh_prev, 4))

One Training Step in Pure NumPy

Before we trust PyTorch's autograd, let us do every operation by hand. The model below is the simplest supervised learner there is — binary logistic regression with two features — but every line maps one-to-one onto what happens inside the LSTM's training step: a forward pass, a loss, an analytic gradient, and a weight update.

         f0      f1
x0     0.50    0.20
x1     0.80    0.40
x2    -0.30   -0.50
x3    -0.60   -0.20

        x0     x1     x2     x3
f0    0.50   0.80  -0.30  -0.60
f1    0.20   0.40  -0.50  -0.20

1import numpy as np
2
3# --- 1. Four labeled training examples (2 features each) ---
4X = np.array([
5    [ 0.5,  0.2],    # "great"   -> y = 1 (positive)
6    [ 0.8,  0.4],    # "amazing" -> y = 1 (positive)
7    [-0.3, -0.5],    # "awful"   -> y = 0 (negative)
8    [-0.6, -0.2],    # "boring"  -> y = 0 (negative)
9])
10y = np.array([1, 1, 0, 0])
11
12# --- 2. Model parameters (logistic regression) ---
13w = np.array([0.1, -0.1])     # starting weights
14b = 0.0                        # starting bias
15lr = 0.5                       # learning rate
16
17def sigmoid(z: np.ndarray) -> np.ndarray:
18    """sigma(z) = 1 / (1 + e^-z) - squashes any real number into (0, 1)."""
19    return 1.0 / (1.0 + np.exp(-z))
20
21# --- 3. FORWARD PASS ---
22z = X @ w + b                  # (4,) logits
23p = sigmoid(z)                 # (4,) probabilities of class 1
24
25# Binary cross-entropy loss (averaged over the batch)
26eps = 1e-9
27loss = -np.mean(y * np.log(p + eps) + (1 - y) * np.log(1 - p + eps))
28
29# --- 4. BACKWARD PASS (analytic gradient of BCE + sigmoid) ---
30# For BCE + sigmoid, dL/dz collapses to (p - y) / n
31grad_z = (p - y) / len(y)      # (4,)
32grad_w = X.T @ grad_z          # (2,)
33grad_b = grad_z.sum()          # scalar
34
35# --- 5. PARAMETER UPDATE (one SGD step) ---
36w = w - lr * grad_w
37b = b - lr * grad_b
38
39print("loss    :", round(loss, 4))
40print("new w   :", np.round(w, 4))
41print("new b   :", round(b, 6))

Click any line of code and the left panel shows the state at that moment — input values, intermediate products, shapes, and the reason every argument is what it is. Notice how the BCE loss of 0.6820 is barely below the random-guess loss of $\ln 2 \approx 0.693$. After the update it would be about 0.660 — a tiny improvement, but repeat this step a few hundred times and the model locks onto the boundary.

The Full Training Loop in PyTorch

PyTorch replaces the NumPy gradient derivation with loss.backward() and the SGD update with optimizer.step(). Everything else — forward pass, loss, the inner loop over batches — looks identical. The full, production-grade template is below.

1import torch
2import torch.nn as nn
3from torch.utils.data import DataLoader
4
5# --- 1. Build the bidirectional LSTM classifier from Section 18.2 ---
6class SentimentLSTM(nn.Module):
7    def __init__(self, vocab_size=5000, d_embed=128, d_hidden=256, n_classes=2):
8        super().__init__()
9        self.embed = nn.Embedding(vocab_size, d_embed, padding_idx=0)
10        self.lstm  = nn.LSTM(d_embed, d_hidden, batch_first=True, bidirectional=True)
11        self.drop  = nn.Dropout(0.4)
12        self.head  = nn.Linear(2 * d_hidden, n_classes)
13
14    def forward(self, x):
15        e      = self.embed(x)                 # (B, T, d_embed)
16        out, _ = self.lstm(e)                  # (B, T, 2*d_hidden)
17        pool   = out.mean(dim=1)               # mean-pool over time
18        return self.head(self.drop(pool))      # (B, n_classes)
19
20# --- 2. Instantiate model, loss, optimizer ---
21model     = SentimentLSTM()
22criterion = nn.CrossEntropyLoss()
23optim     = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-5)
24
25# --- 3. Data loaders built earlier from tokenised reviews ---
26train_loader = DataLoader(train_ds, batch_size=64, shuffle=True)
27val_loader   = DataLoader(val_ds,   batch_size=256)
28
29# --- 4. Training loop with early stopping ---
30best_val, bad, PATIENCE = float("inf"), 0, 4
31
32for epoch in range(1, 26):
33    # ---- train ----
34    model.train()
35    train_loss = 0.0
36    for xb, yb in train_loader:
37        optim.zero_grad()                      # wipe previous gradients
38        logits = model(xb)                     # forward
39        loss   = criterion(logits, yb)         # scalar tensor
40        loss.backward()                        # autograd
41        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
42        optim.step()                           # Adam update
43        train_loss += loss.item() * xb.size(0)
44    train_loss /= len(train_ds)
45
46    # ---- validate ----
47    model.eval()
48    val_loss, correct = 0.0, 0
49    with torch.no_grad():
50        for xb, yb in val_loader:
51            logits = model(xb)
52            val_loss += criterion(logits, yb).item() * xb.size(0)
53            correct  += (logits.argmax(dim=1) == yb).sum().item()
54    val_loss /= len(val_ds)
55    val_acc   = correct / len(val_ds)
56    print(f"ep {epoch:2d} | train {train_loss:.3f} | val {val_loss:.3f} | acc {val_acc:.3f}")
57
58    # ---- early stopping ----
59    if val_loss < best_val:
60        best_val, bad = val_loss, 0
61        torch.save(model.state_dict(), "best.pt")
62    else:
63        bad += 1
64        if bad >= PATIENCE:
65            print(f"early stop at epoch {epoch}")
66            break

Five things in that template are worth internalising because you will write them in every training loop for the rest of your career:

1. model.train() / model.eval() — flip dropout and batchnorm on/off.
2. optim.zero_grad() at the top of every step — PyTorch accumulates gradients by default, which is almost never what you want.
3. with torch.no_grad(): around evaluation — disables the computation graph, halves memory, doubles speed.
4. clip_grad_norm_ — a cheap safety belt against exploding gradients that RNNs are notorious for.
5. torch.save(model.state_dict(), ...) on every best-so-far epoch — cheaper than re-training when something goes wrong.

Evaluation: Accuracy, Precision, Recall, F1

Accuracy — the fraction of predictions we got right — is the number you watch first, but it can be misleading. If 95% of your reviews are positive, a model that always predicts "positive" scores 95% accuracy without doing any work. We need richer metrics, and all of them derive from four counts organised in the confusion matrix.

From those four counts we derive:

• Accuracy $= \frac{TP + TN}{TP + TN + FP + FN}$ — fraction correct. Dominated by the majority class in imbalanced datasets.
• Precision $= \frac{TP}{TP + FP}$ — of the reviews we flagged positive, how many actually were? High precision = few false alarms.
• Recall $= \frac{TP}{TP + FN}$ — of the actually-positive reviews, how many did we catch? High recall = few misses.
• F1-score $= \frac{2 \cdot P \cdot R}{P + R}$ — harmonic mean of precision and recall. Punishes models that trade one against the other. The default single-number metric for classification.

Why the harmonic mean? Because if either P or R is close to zero, the harmonic mean is also close to zero. The arithmetic mean $(P + R)/2$ would let a model cheat by being great at one metric and terrible at the other. F1 says you must do well on both.

Threshold matters — precision/recall is a knob

For a binary classifier, the prediction comes from $p(x) > \tau$. Default $\tau = 0.5$, but you can slide it. Raise the threshold and precision goes up (we only call something positive if we are very sure), at the cost of recall. Lower it and recall rises but precision suffers. The right setting depends on which mistake costs more — an FP or an FN.

Computing the Metrics in Code

All four metrics drop out of a few NumPy boolean operations. You will rarely write this by hand in production — sklearn, torchmetrics, and HuggingFace's evaluate all ship one-liners — but understanding the counts behind the numbers is essential.

1import numpy as np
2
3# A model's predictions on a 20-sample validation slice.
4# 1 = positive review, 0 = negative review.
5y_true = np.array([1, 0, 1, 1, 0, 1, 0, 0, 1, 0, 1, 1, 0, 1, 0, 0, 1, 1, 1, 0])
6y_pred = np.array([1, 0, 1, 0, 0, 1, 1, 0, 1, 0, 1, 0, 0, 1, 0, 1, 1, 1, 1, 0])
7
8# --- Confusion matrix counts ---
9tp = int(((y_pred == 1) & (y_true == 1)).sum())   # predicted pos, actual pos
10fp = int(((y_pred == 1) & (y_true == 0)).sum())   # predicted pos, actual neg
11fn = int(((y_pred == 0) & (y_true == 1)).sum())   # predicted neg, actual pos
12tn = int(((y_pred == 0) & (y_true == 0)).sum())   # predicted neg, actual neg
13
14# --- Derived metrics ---
15accuracy  = (tp + tn) / (tp + tn + fp + fn)
16precision = tp / max(1, tp + fp)
17recall    = tp / max(1, tp + fn)
18f1        = 2 * precision * recall / max(1e-9, precision + recall)
19
20print(f"TP={tp}  FP={fp}  FN={fn}  TN={tn}")
21print(f"accuracy={accuracy:.3f}  precision={precision:.3f}  "
22      f"recall={recall:.3f}  f1={f1:.3f}")

Plug in one of the epochs from the learning-curve chart below. At epoch 13 the validation set had TP=139, FN=9, FP=72, TN=280, giving $\text{accuracy} \approx 0.838$, $P \approx 0.659$, $R \approx 0.939$, and $F_1 \approx 0.774$. The precision is dragged down by a large number of false positives — a signal that a higher decision threshold would trade a little recall for a lot of precision.

The knob matters. Below, drag the decision threshold $\tau$ and watch precision climb as you demand more confidence — at the cost of recall. The PR curve traces every operating point your classifier can reach without retraining.

Learning Curves and the Overfitting Story

Plotting train loss, val loss and accuracy against epoch is the single most important diagnostic you can run on a model. Three failure modes reveal themselves here before they poison your test set.

Drag the epoch slider across a realistic 25-epoch run of this exact LSTM classifier. Read the curves this way:

1. Epochs 1-6 — learning. Both losses drop quickly. Train and val curves hug each other. The model is discovering the sentiment signal.
2. Epochs 7-13 — consolidation. Rate of improvement slows. Val loss is still decreasing. This is your "sweet spot" region, and the green best line marks its minimum.
3. Epochs 14-25 — overfitting. Train loss keeps dropping, but val loss starts rising and val accuracy stalls or decays. The model is memorising idiosyncrasies of the training set — rare reviewer phrases, punctuation patterns — that do not generalise.

Underfitting shows up as the opposite pattern: both losses stay high and the curves plateau early. Remedy: make the model bigger, lower regularization, train longer, or feed richer features.

Regularization: Dropout, Weight Decay, Early Stopping

Overfitting is the phenomenon that the model's capacity exceeds what the data actually supports. Every practical regularization technique is a constraint that shrinks effective capacity without shrinking raw parameter count.

Dropout — the Bayesian intuition

Dropout can be read as training an exponentially large ensemble of thinned sub-networks that share weights. With p = 0.4 on a 512-dim pooled vector, there are $2^{512}$ possible masks. At test time we average by leaving all units on and scaling — a free approximation to the ensemble's mean prediction. The ensemble lens explains why dropout reduces variance without hurting expressivity.

Dropout was introduced by Srivastava et al. 2014 (JMLR 15) and first tuned for RNN-specific use by Zaremba, Sutskever & Vinyals 2014 (arXiv:1409.2329), which is why our RNN dropout rates (0.3–0.5) sit higher than the transformer-era defaults (0.1–0.3).

Weight decay — keep the weights small

Adding $\tfrac{\lambda}{2} \|\theta\|^2$ to the loss and differentiating gives $\theta \leftarrow \theta - \eta(g + \lambda \theta) = (1 - \eta\lambda)\theta - \eta g$. Every step multiplies the weight by $(1 - \eta\lambda) < 1$, pulling it slightly toward zero. Large weights need a reason to survive.

Early stopping — the free regularizer

Early stopping is often the simplest regularizer to deploy. It costs nothing at train time, has one hyperparameter (patience), and works regardless of architecture. The caveat: you need a validation set you trust.

Scaling the Same Recipe to Transformers

Everything we have written for the LSTM — train mode, mini-batch loop, cross-entropy, Adam, early stopping — transfers to transformers essentially unchanged. What changes is scale and the bottlenecks that scale exposes. Below is the training anatomy of a modern transformer language model at roughly 1 B parameters, side-by-side with our ~1.8 M-parameter LSTM.

Every line of that table is the same five-step training loop wearing a different coat. A few extensions deserve mention because they are direct consequences of scaling the recipe:

1. Learning-rate schedule. Transformers warmup $\eta$ linearly for the first 1-4 k steps (because $\hat{v}$ needs time to converge to a good estimate), then cosine-decay it over the remainder of training. This came out of the original transformer paper and survives to this day.
2. Gradient accumulation. When the desired batch (e.g., 4 M tokens) does not fit in GPU memory, weaccumulate gradients over several micro-batches before calling optimizer.step(). Same math, lower memory peak.
3. Mixed-precision training. Forward and backward in bf16 or fp8, but Adam's moments in fp32. Roughly halves memory and 2x speed on A100/H100, with no accuracy penalty if done correctly.
4. Parallelism. Data parallel across GPUs, tensor parallel across GPUs, pipeline parallel across GPUs, and ZeRO sharding across GPUs — all orthogonal strategies to scale one training step across a cluster.

Flash Attention: Fixing the Training-Time Memory Wall

If you take our five-step loop and naively swap the LSTM for a transformer, the first wall you hit is not compute — it is activation memory. Self-attention produces a score matrix of shape $(T, T)$ at every layer. For $T = 8\,192$ that is $8192^2 = 6.7 \times 10^{7}$ floats per head, and we need to keep it around for the backward pass. A 32-layer, 32-head model at 8 k context chews through hundreds of gigabytes of activation memory before anyone trains anything.

Flash Attention (Dao et al., 2022) solves this by rewriting self-attention as a tiled, online softmax. It never materialises the full $(T, T)$ score matrix in high-bandwidth memory (HBM). Instead it:

1. Loads small tiles of $Q, K, V$ into SRAM.
2. Computes each tile's contribution to the softmax on the fly, updating a running maximum and running sum using the numerically-stable online softmax algorithm.
3. Writes only the attention output back to HBM — which is $(T, d)$, not $(T, T)$.
4. For the backward pass, recomputes the tiles on the fly rather than caching them — trading cheap SRAM compute for expensive HBM storage.

Flash Attention is the single most important implementation-level paper of the transformer era. It is not a new model, a new loss, or a new algorithm — it is a memory-aware reorganisation of the exact math we wrote above. Yet it unlocked long-context training on consumer-grade accelerators and is the reason modern LLMs reason over book-length contexts during training.

The LSTM classifier in this chapter never needs any of this. Its hidden-state tensor at one timestep is a single $d_{\text{hidden}}$-dim vector, not a quadratic score matrix. This is the memory upside of recurrence — cheap in VRAM, expensive in wall-clock time. Flash Attention is the engineering answer to the question "how do we keep the matrix-multiply speed of a transformer while giving back the memory efficiency of an LSTM?"

KV-Cache: Why Evaluation Looks Different from Training

There is one more place where our LSTM and a transformer diverge sharply, and it shows up at evaluation / inference time — the same torch.no_grad() block in our training loop.

During LSTM eval, decoding a new token means one cheap gate computation that consumes the previous hidden and cell state and produces the next. The only "cache" we carry forward is those two $d_{\text{hidden}}$-dim vectors.

During transformer eval, generating token $T+1$ means re-running self-attention where the new token's query attends to all previous tokens' keys and values. Recomputing those keys and values from scratch at every step would make generation $O(T^2)$ per token.

The KV-cache saves K and V matrices for every past token, layer-by-layer, so each new step only computes one new Q/K/V triple and attends against the growing cache. The price is memory:

At 70 B scale the KV-cache alone can exceed the model weights in memory. This drove the invention of several variants:

• Multi-Query Attention (MQA) — share one K/V head across all Q heads. Cuts KV-cache size by $h$ (head count), at a small quality cost.
• Grouped-Query Attention (GQA) — compromise: $g$ K/V heads shared across $h$ Q heads (Llama 2/3 use $g = 8$).
• Paged Attention (vLLM) — the OS's virtual memory trick applied to attention: the KV-cache is stored in fixed-size pages so batches of different sequence lengths don't waste space.

Summary

In this section we closed the loop on the sentiment classifier we have been building for three sections:

1. Stated training as one optimization problem over parameters and showed that the same equation drives everything from logistic regression to a 1 B-parameter transformer.
2. Derived cross-entropy from three angles (likelihood, information theory, gradient friendliness) and showed why it pairs so cleanly with sigmoid / softmax.
3. Walked from SGD to momentum to Adam, visualized the difference with an interactive 3D loss surface, and explained why Adam (and AdamW) is the default optimizer of modern deep learning.
4. Traced one complete training step in pure NumPy — forward, loss, backward, update — with every numeric value computed. Then wrote the same recipe as an industrial PyTorch training loop, annotating every function call and argument.
5. Defined accuracy, precision, recall, F1 and the confusion matrix from first principles, and scrubbed through a realistic 25-epoch run to see overfitting appear exactly where the theory predicted.
6. Cataloged the regularization toolkit — dropout, weight decay, early stopping, gradient clipping, augmentation — and showed which one each hyperparameter in our training loop corresponds to.
7. Extended the same five-step loop to billion-parameter transformers and identified what scaling breaks: activation memory (Flash Attention fixes it) and inference memory (the KV-cache and its MQA / GQA / paged variants fix it).

The line from our modest LSTM to modern LLMs is remarkably straight. Swap the architecture, swap the dataset, swap the hardware — but the training loop, the loss, the optimizer, the metrics and even the overfitting diagnostic are the same five steps we just wrote. Everything else is engineering in service of making those five steps fit onto the hardware.

References

• Kingma, D. P. & Ba, J. (2014). Adam: A Method for Stochastic Optimization. arXiv:1412.6980.
• Srivastava, N., Hinton, G., Krizhevsky, A., Sutskever, I. & Salakhutdinov, R. (2014). Dropout: A Simple Way to Prevent Neural Networks from Overfitting. JMLR 15.
• Zaremba, W., Sutskever, I. & Vinyals, O. (2014). Recurrent Neural Network Regularization. arXiv:1409.2329.
• Loshchilov, I. & Hutter, F. (2019). Decoupled Weight Decay Regularization. ICLR 2019 / arXiv:1711.05101.
• Dao, T., Fu, D. Y., Ermon, S., Rudra, A. & Ré, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022 / arXiv:2205.14135.
• Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebrón, F. & Sanghai, S. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP 2023 / arXiv:2305.13245.
• Kwon, W. et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. SOSP 2023 / arXiv:2309.06180.
• Grattafiori, A. et al. (2024). The Llama 3 Herd of Models. arXiv:2407.21783.