DL S011 · MLP on MNIST — All NumPy, No PyTorch

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 11 / 80 · Module M02 · ~2.5 hours

0 · The story — a French scientist, a cheque-reading machine, and the 200/second rule

1998. Yann LeCun is at AT&T Bell Labs in Holmdel, New Jersey, working on a machine to read the amount fields on paper cheques. The system is called NCR-5000. By the end of the year it will be reading, at production banks, roughly 10% of all cheques written in the United States — some 200 cheques per second per machine. The neural network at the core of that system is LeNet-5: a convolutional network trained on a dataset that LeCun and his colleagues put together from a mix of NIST hand-written character databases (from prison-inmate and U.S. Census questionnaires) and re-standardised. They called it the Modified NIST, or MNIST.

MNIST is 60,000 training images and 10,000 test images of hand-written digits, each 28×28 greyscale, labelled 0–9. Small enough to fit in RAM as a NumPy array. Simple enough that an MLP works. Famous enough that when you hit 97% today, you have — in ~50 lines of pure NumPy on a laptop — approximately matched LeCun's original 1998 LeNet-5 accuracy on the same dataset. LeCun needed a Sun SPARCstation. You need python -c "...".

Every ML engineer has a story about the first neural network they trained. Mine was in 2018, in TensorFlow 1.x with placeholders and sessions and a graph I didn't understand. It worked and I did not know why. That's a bad way to learn. This session gives you the opposite experience: every line yours, every gradient defensible.

MNIST is also — and this is under-appreciated — still an active research probe in 2025. When researchers want to understand phenomena that only appear in training dynamics (grokking, phase transitions, linear mode connectivity, lottery tickets) they usually re-run on MNIST-scale MLPs because a full ablation of a phenomenon on Llama-3-70B costs millions of dollars, but on your MLP it costs seconds. Nanda, Chan, Lieberum, Smith & Steinhardt's "Progress measures for grokking via mechanistic interpretability" (ICLR 2023) ran hundreds of ablations on a two-layer MLP. Frankle & Carbin's "The Lottery Ticket Hypothesis" (ICLR 2019) started on MNIST. MNIST is not obsolete; it is a scientific instrument.

Today: build the instrument.

1 · What you'll build today

A 48-line NumPy script that loads MNIST, defines a 784-128-10 MLP with softmax + CE and He init, trains it with mini-batch SGD for 10 epochs, and reports >97% test accuracy.

2 · Checkpoint · you should already know

• S005 · Probability & information — softmax + cross-entropy, and its clean P − y_one_hot gradient (§7 of S005).
• S006 · Optimisation intuition — SGD, mini-batches, and the log-sum-exp trick.
• S009 · MLP forward — the 2-layer forward pass on tiny numbers.
• S010 · Backprop by hand — the 15-line backward pass and the cheat-sheet formulas.

If any of those feel foggy, a 10-minute skim of S010 §4 (the cheat-sheet table) is by far the highest-leverage prep.

3 · Load MNIST — no framework, just NumPy

We won't use torchvision or Keras — we want raw NumPy. The easiest source is sklearn.datasets.fetch_openml (one line, one download, cached on disk after):

from sklearn.datasets import fetch_openml
import numpy as np
 
mnist = fetch_openml("mnist_784", version=1, as_frame=False, cache=True)
X, y = mnist.data.astype(np.float32), mnist.target.astype(np.int64)
 
# Normalise pixels to [0, 1]
X /= 255.0
 
# Split — the classic 60k train / 10k test
X_train, X_test = X[:60000], X[60000:]
y_train, y_test = y[:60000], y[60000:]
 
print(X_train.shape, y_train.shape, X_test.shape, y_test.shape)
# (60000, 784) (60000,) (10000, 784) (10000,)

Each row of X is a flattened 28×28 image (784 pixels). Each label is 0–9. That is the whole dataset. Total memory: ~180 MB in float32. Fits comfortably in RAM on any laptop made after ~2005.

3.1 Visualise a few — always look at your data

import matplotlib.pyplot as plt
fig, axes = plt.subplots(1, 8, figsize=(12, 2))
for i, ax in enumerate(axes):
    ax.imshow(X_train[i].reshape(28, 28), cmap="gray")
    ax.set_title(int(y_train[i])); ax.axis("off")
plt.show()

Squint at them. Note the ones that are visually ambiguous (a 4 that looks like a 9, a scribbly 5). Even a perfect model won't get 100% because some labels are wrong or ambiguous. The ceiling on MNIST accuracy is estimated at ~99.7%; anything above is probably fitting label noise. Karpathy's MNIST comparison notebook is a good rabbit hole.

3.2 Sanity checks on the loaded arrays

print("X_train mean:", X_train.mean(), "  std:", X_train.std())
# X_train mean: ~0.13    std: ~0.31    (roughly — 87% of pixels are black)
 
print("Class balance (train):", np.bincount(y_train))
# [5923 6742 5958 6131 5842 5421 5918 6265 5851 5949]   ≈ balanced

Class balance matters — an unbalanced dataset needs different loss weighting (S065). MNIST is nice: within ~10% across classes.

4 · The 784-128-10 MLP with softmax

Architecture:

input:  x ∈ ℝ^{784}
Layer 1: h  = ReLU(x @ W1 + b1),   W1: (784, 128),  b1: (128,)
Layer 2: z  = h @ W2 + b2,         W2: (128, 10),   b2: (10,)
output:  p  = softmax(z)           probability over 10 classes
loss:    L  = -log p_y             averaged over batch

Parameter count: (784·128 + 128·10) + (128 + 10) = 100,352 + 1,280 + 138 = **101,770**. About 100k parameters. Tiny by modern standards (a single attention head in Llama-3-70B has ~4M); huge for a Sun SPARCstation in 1998.

4.1 He initialisation — derived, not memorised

Kaiming He's 2015 paper "Delving deep into rectifiers" (arxiv 1502.01852) noticed that for ReLU layers, if you want the variance of the activations to be preserved through depth, you need weights drawn from:

Where does the 2/n_in come from? For a linear layer z_j = Σᵢ Wᵢⱼ xᵢ, if x and W are independent with mean 0 and variances σ_x², σ_W², then Var(z_j) = n_in · σ_W² · σ_x². For variance preservation you'd set σ_W² = 1/n_in (Xavier init, Glorot & Bengio 2010). But ReLU zeros out half the activations, halving their variance downstream. To compensate, double the weight variance: σ_W² = 2/n_in. That's it. That's He init in one paragraph.

def init_params(sizes=(784, 128, 10), seed=0):
    rng = np.random.default_rng(seed)
    W1 = rng.standard_normal((sizes[0], sizes[1])) * np.sqrt(2.0 / sizes[0])
    b1 = np.zeros(sizes[1])
    W2 = rng.standard_normal((sizes[1], sizes[2])) * np.sqrt(2.0 / sizes[1])
    b2 = np.zeros(sizes[2])
    return W1, b1, W2, b2

Bias zero, weights He-scaled. Boring, correct, and enough. S022 covers the full init family (Xavier, He, LSUV, orthogonal, μP) in depth.

5 · Forward pass — 9 lines

def softmax_stable(z):
    # log-sum-exp trick (S006 §6.2): subtract max for numerical stability
    z = z - z.max(axis=1, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=1, keepdims=True)
 
def forward(X, params):
    W1, b1, W2, b2 = params
    Z1 = X @ W1 + b1                              # (N, 128)
    H  = np.maximum(0, Z1)                        # (N, 128)   ReLU
    Z2 = H @ W2 + b2                              # (N, 10)
    P  = softmax_stable(Z2)                       # (N, 10)
    return P, (X, Z1, H, Z2, P)

The softmax_stable shift is not decoration. Without it, a logit of 100 → exp(100) ≈ 2.7·10⁴³ → overflows float32 to inf → inf / inf → NaN → training dies at epoch 1. Always shift.

6 · Loss and backward — 15 lines

def cross_entropy(P, y):
    N = y.size
    return -np.log(P[np.arange(N), y] + 1e-12).mean()
 
def backward(y, cache, params):
    X, Z1, H, Z2, P = cache
    W1, b1, W2, b2 = params
    N = y.size
 
    # softmax + CE — the clean fused gradient (S005 §7): dZ2 = P - one_hot(y), /N
    dZ2 = P.copy()
    dZ2[np.arange(N), y] -= 1
    dZ2 /= N                                       # (N, 10)
 
    # layer 2 backward (cheat-sheet from S010)
    dW2 = H.T @ dZ2                                # (128, 10)
    db2 = dZ2.sum(axis=0)                          # (10,)
 
    # ReLU backward
    dH  = dZ2 @ W2.T                               # (N, 128)
    dZ1 = dH * (Z1 > 0)                            # (N, 128)
 
    # layer 1 backward
    dW1 = X.T @ dZ1                                # (784, 128)
    db1 = dZ1.sum(axis=0)                          # (128,)
 
    return dW1, db1, dW2, db2

Fifteen lines. Exactly the pattern from S010 §7, with softmax + CE instead of sigmoid + BCE. Same shape rules, same cheat-sheet formulas, wider tensors. The trick dZ2[np.arange(N), y] -= 1 performs P − one_hot(y) without materialising the one-hot matrix — a 10× memory saving on this layer and an idiom you will see in every PyTorch loss implementation.

7 · Training loop — 20 lines with mini-batches

def train(X, y, X_test, y_test, epochs=10, batch_size=64, lr=0.1, seed=0):
    rng = np.random.default_rng(seed)
    params = init_params(seed=seed)
    n = X.shape[0]
    for epoch in range(epochs):
        # shuffle — MANDATORY (see war story #1)
        idx = rng.permutation(n)
        X_shuf, y_shuf = X[idx], y[idx]
        for start in range(0, n, batch_size):
            xb = X_shuf[start:start+batch_size]
            yb = y_shuf[start:start+batch_size]
            _, cache = forward(xb, params)
            dW1, db1, dW2, db2 = backward(yb, cache, params)
            W1, b1, W2, b2 = params
            params = (W1 - lr*dW1, b1 - lr*db1, W2 - lr*dW2, b2 - lr*db2)
        # eval
        P_test, _ = forward(X_test, params)
        acc = (P_test.argmax(1) == y_test).mean()
        loss = cross_entropy(P_test, y_test)
        print(f"epoch {epoch+1:2d}: test acc = {acc:.4f}  test loss = {loss:.4f}")
    return params

Let's read the loop end-to-end.

• Shuffle every epoch → each mini-batch is a fresh random sample. This is what the "S" in SGD means.
• Batch size 64 → each gradient step uses 64 examples. Big enough for a stable gradient estimate, small enough that noise helps escape saddle points and shallow minima (S006 §2.1). The linear scaling rule (Goyal et al. 2017) says if you multiply batch by k, multiply lr by k (up to a point).
• lr = 0.1 → moderate for MNIST + He init + softmax-CE. Try 0.01 for a smoother curve; try 1.0 to see divergence.
• 10 epochs → 10 passes over the training set = 10 · 60000 / 64 ≈ 9375 gradient steps total. Enough for near-convergence on this problem.

Run it:

params = train(X_train, y_train, X_test, y_test, epochs=10, batch_size=64, lr=0.1)
# epoch  1: test acc = 0.9412  test loss = 0.2017
# epoch  2: test acc = 0.9581  test loss = 0.1445
# epoch  3: test acc = 0.9642  test loss = 0.1229
# ...
# epoch 10: test acc = 0.9756  test loss = 0.0842

97.5% test accuracy. In fewer than 50 lines of NumPy. On CPU. In under 90 seconds. LeCun's 1998 LeNet-5 (a convnet, more sophisticated than what we built) hit ~99.2% — but he needed a Sun SPARCstation and weeks of training. You just spent 90 seconds on hardware LeCun would have called science-fiction.

8 · Read the loss curve — the deep-learning story in one plot

Plot the per-epoch loss. It should go from ~2.3 (= ln 10, uniform-random on 10 classes) down toward ~0.08. That's the entire deep-learning story: define a loss, follow the gradient, get better. Every model you'll ever train produces a curve like this — for MNIST it drops fast; for a 70B LLM it drops for months.

Reading the curve like an X-ray:

• Starting loss is ln K for a K-way classifier with uniform init on the output. ln 10 ≈ 2.303. If step 0 isn't near this, your output layer's init is asymmetric or your softmax has a bug.
• Steady decay, then plateau → converged. Fine.
• Curve gets noisier as it goes down → gradient variance is high relative to loss magnitude; consider a smaller lr near the end or an lr schedule (S023).
• Spikes → lr too high; a single bad batch is enough to jump uphill.
• Sudden jump to NaN → softmax overflow (revisit §5) or gradient explosion (S024 talks about clipping).
• Train loss keeps decreasing but test loss increases → overfitting (dropout in S020 will help).

9 · Seven ablations — 30 minutes very well spent

Run these variations. Notice what changes.

1. lr = 1.0 instead of 0.1 → loss spikes; may diverge; if it converges, final accuracy typically worse. Learning rate is the single most important hyperparameter.
2. batch_size = 1 (pure SGD) → very noisy loss, but often finds a slightly better minimum (Keskar et al. 2016 argued small-batch minima generalise better). Ten times slower per epoch though — fewer vectorised operations.
3. batch_size = 1024 → smoother loss curve, but takes ~5× longer per epoch to reach the same accuracy — fewer gradient steps per epoch dominates. Also: large-batch training often needs lr warmup (S023).
4. Remove He init (np.random.randn * 0.01) → training 2–3× slower to converge; final accuracy 1–2% worse. He init is not decoration.
5. Remove ReLU (identity everywhere) → the whole network collapses to a linear model (S008 §3). Test accuracy plateaus at ~92%. Surprisingly good — MNIST is almost linear! This is a famous MNIST oddity; it's why serious vision benchmarks moved to CIFAR-10/100 and ImageNet.
6. Add a third layer (784-128-64-10) → about 0.5% better, 1.5× compute. Diminishing returns on MNIST.
7. Shuffle disabled → accuracy plateaus at ~89% (see war story #1). MNIST's default label ordering (0000...111...) means every mini-batch is one class → biased gradients.

Do at least four of these. Every one builds intuition you'll use for the rest of your ML career.

10 · The complete script — 48 lines to memorise

import numpy as np
from sklearn.datasets import fetch_openml
 
mnist = fetch_openml("mnist_784", version=1, as_frame=False, cache=True)
X = mnist.data.astype(np.float32) / 255.0
y = mnist.target.astype(np.int64)
X_train, X_test = X[:60000], X[60000:]
y_train, y_test = y[:60000], y[60000:]
 
def init(sizes=(784, 128, 10)):
    return (np.random.randn(sizes[0], sizes[1]) * np.sqrt(2/sizes[0]),
            np.zeros(sizes[1]),
            np.random.randn(sizes[1], sizes[2]) * np.sqrt(2/sizes[1]),
            np.zeros(sizes[2]))
 
def softmax(z):
    z = z - z.max(axis=1, keepdims=True)
    e = np.exp(z);  return e / e.sum(axis=1, keepdims=True)
 
def forward(X, params):
    W1, b1, W2, b2 = params
    Z1 = X @ W1 + b1
    H  = np.maximum(0, Z1)
    Z2 = H @ W2 + b2
    P  = softmax(Z2)
    return P, (X, Z1, H, Z2)
 
def backward(y, cache, params):
    X, Z1, H, Z2 = cache
    W1, b1, W2, b2 = params
    N = y.size
    dZ2 = softmax(Z2).copy();  dZ2[np.arange(N), y] -= 1;  dZ2 /= N
    dW2 = H.T @ dZ2;    db2 = dZ2.sum(0)
    dH  = dZ2 @ W2.T
    dZ1 = dH * (Z1 > 0)
    dW1 = X.T @ dZ1;    db1 = dZ1.sum(0)
    return dW1, db1, dW2, db2
 
def train(X, y, Xt, yt, epochs=10, bs=64, lr=0.1):
    p = init()
    n = X.shape[0]
    for e in range(epochs):
        idx = np.random.permutation(n)
        Xs, ys = X[idx], y[idx]
        for s in range(0, n, bs):
            xb, yb = Xs[s:s+bs], ys[s:s+bs]
            _, cache = forward(xb, p)
            g = backward(yb, cache, p)
            p = tuple(pi - lr*gi for pi, gi in zip(p, g))
        Pt, _ = forward(Xt, p)
        print(f"epoch {e+1}: acc={(Pt.argmax(1)==yt).mean():.4f}")
    return p
 
train(X_train, y_train, X_test, y_test)

Forty-eight lines. 97% MNIST accuracy. Print it. Frame it. Bring it into your next ML interview.

11 · Why MNIST still matters in 2025 — three research examples

MNIST is not obsolete; it is a scientific instrument. Three examples of frontier work in the last few years where MNIST-scale MLPs are the main probe:

1. Grokking (Power et al. 2022, arxiv 2201.02177 → follow-up mechanistic-interpretability work Nanda et al. 2023, arxiv 2301.05217). Networks trained on tiny algorithmic tasks memorise the training data quickly, then — thousands of steps after memorisation — suddenly generalise. The mechanism is a phase transition in the network's internal representations. Studying it requires re-running training thousands of times, which is only feasible on MNIST-scale models.
2. Lottery Ticket Hypothesis (Frankle & Carbin, ICLR 2019, arxiv 1803.03635). Started on MNIST + LeNet, then Fashion-MNIST + CIFAR. Claim: inside a dense network, there exists a sparse "winning ticket" subnetwork that trains to full accuracy on its own. Sparked a huge sub-field of pruning research.
3. Linear Mode Connectivity (Frankle et al. 2020, arxiv 1912.05671). Two independently trained MNIST MLPs from the same init end up in the same loss basin — you can linearly interpolate their weights and loss stays low. Doesn't hold for random inits. This is what makes model soups (Wortsman et al. 2022) and weight averaging work in modern LLM training.

Further reading:

• LeCun et al., "Gradient-Based Learning Applied to Document Recognition" (1998) — http://yann.lecun.com/exdb/publis/pdf/lecun-01a.pdf
• He et al., "Delving Deep into Rectifiers" (He init, 2015) — https://arxiv.org/abs/1502.01852
• Power et al., "Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets" (2022) — https://arxiv.org/abs/2201.02177
• Nanda et al., "Progress measures for grokking via mechanistic interpretability" (ICLR 2023) — https://arxiv.org/abs/2301.05217
• Frankle & Carbin, "The Lottery Ticket Hypothesis" (ICLR 2019) — https://arxiv.org/abs/1803.03635

12 · War stories

13 · Try it yourself

14 · Recap

You wrote every line: data load, He-init, softmax + CE forward, P − one_hot(y) fused backward, mini-batch SGD with shuffling. 48 lines total, 97% test accuracy on the dataset LeCun cut deep learning's teeth on. You matched a 1998 result on 2025 hardware in under 90 seconds. Along the way you saw the sister of a − y (namely P − one_hot(y)), derived He init from a variance argument, and heard how MNIST is still the scientific instrument of choice for phenomena like grokking and lottery tickets. Everything from M01 and M02 has cashed out into one script. Congratulations.

🧠 Retention scaffold

One-line summary (write it in your own words): _______________________________________________________________

Spaced review: re-run the 48-line script from memory in 24 hours (no copy-paste); revisit §9 (ablations) on day 7 with different hyperparameters.

Next session (S012): Karpathy's micrograd — a ~100-line pure-Python autograd engine built around a Value class. You'll build the mathematical seed of what PyTorch does with tensors, and gain the "aha" that turns autograd from magic into 40 lines of code you can hold in your head.

Sticky note (keep on your desk):

• 48-line script. Print it. Really.
• Fact: 100k params + He init + softmax-CE + SGD lr=0.1 + shuffle → 97% MNIST. Reproducible on any laptop in 90s.
• Snippets: dZ2[np.arange(N), y] -= 1; dZ2 /= N — softmax+CE grad in two lines. W ~ N(0, sqrt(2/n_in)) — He init in one line.
• Reading rule: step-0 loss for K-way softmax on balanced data is ln K. Memorise ln 10 ≈ 2.303.

Previous: ← DL S010 · Backprop by Hand · Next: DL S012 · Micrograd — Build a Tiny Autograd Engine →