mnist_nn_scratch

A neural network built from NumPy and nothing else — forward pass, backprop, convolutions, optimizers, batch norm, dropout, all hand-written. No PyTorch, no TensorFlow, no autograd.

99.39% test accuracy on MNIST (61 errors / 10,000) from the conv net with augmentation, 99.55% with test-time augmentation on top (averaged over three seeds; any single run lands within about two errors of that), or 98.28% from the MLP in 60 seconds.

Quickstart

python mnist_nn_scratch/gradcheck.py

python mnist_nn_scratch/train.py --arch cnn --augment --epochs 15 --plot

python mnist_nn_scratch/evaluate.py --errors-png

That is the headline run, half an hour or so on a laptop CPU. Drop --augment --epochs 15 for the 13-minute version at 99.28%, or use --arch mlp to see the whole thing work end to end in a minute.

train.py downloads MNIST (11.5 MB) on first run and caches it under data/. Requires Python 3 and NumPy; matplotlib only for the optional plots.

Results

	params	epochs	time	test accuracy	errors
--arch mlp	235,530	20	60 s	98.28%	172
--arch mlp --augment	235,530	30	3.7 min	98.93%	107
--arch cnn	215,546	8	13 min	99.28%	72
--arch cnn --augment	215,546	15	56 min*	99.39%	61
--arch cnn --augment --elastic	215,546	15	103 min*	98.62%	138

* measured wall time, but both of those runs shared the machine, so read them as upper bounds rather than costs — per-epoch time ranged 105–375 s for the augmented run and 285–658 s for the elastic one, a spread that is contention, not the workload. Times here are single measurements on one laptop CPU, not benchmarks. The warping itself is cheap next to the network: over a 55,000 image epoch, affine costs 16 s and affine+elastic 30 s, against roughly 98 s for an unaugmented conv epoch.

The last row is a negative result kept on the record deliberately: elastic distortion, the one addition here that a paper would tell you should help, made things worse. Why, and how confidently, is in the elastic entry under Decisions.

Two independent effects, and they compose. Convolution cuts the MLP's error rate by 58% using fewer parameters — weight sharing beats brute-force width, which is the whole argument for convolution. Augmentation then cuts what is left by a further 15%, and it costs no parameters at all: the same model, shown warped copies of the same 55,000 images.

Augmentation is worth far more to the MLP (38% of its errors) than to the conv net (15%). That ordering makes sense — small shifts and rotations are exactly the invariance a conv net already has structurally, so teaching it by example is partly redundant, whereas the MLP has no such prior and has to learn every translation from data.

Test-time augmentation

The same warps applied at inference, averaging predictions over the clean image plus N warped views. No retraining — it buys accuracy for inference time.

python mnist_nn_scratch/evaluate.py --checkpoint checkpoints/cnn_aug.npz --tta 8

Errors out of 10,000, mean ± spread over seeds 0/1/2:

views	mlp+aug (107 base)	cnn+aug (61 base)
N=1	98.3 ± 1.2	50.0 ± 4.5
N=2	97.7 ± 0.9	47.7 ± 2.5
N=4	93.0 ± 3.7	48.3 ± 2.5
N=8	91.7 ± 2.4	45.3 ± 1.7
N=16	93.3 ± 2.1	46.0 ± 0.8

That is 61 → 45 errors on the conv net, 99.39% → 99.55%, and 107 → 92 on the MLP, 98.93% → 99.08%. Both gains are around 15 errors against a seed-to-seed spread of 2, so they are real and not a lucky draw — which is the only reason this table is worth printing. The curve flattens by N=8; the spread keeps shrinking after that, but the mean does not improve, so more views buy consistency rather than accuracy.

MLP — 784 → 256 → 128 → 10, ReLU, batch norm, 20% dropout, Adam with cosine decay and warmup.

epoch   1/20  loss 0.5257  train 84.69%  val 95.30%
epoch  20/20  loss 0.0152  train 99.56%  val 98.16%
best val 98.18% at epoch 19 -> test 98.28%  (172 errors)

CNN — Conv(1→16, 5×5) → BN → ReLU → pool → Conv(16→32, 5×5) → BN → ReLU → pool → Dense(1568→128) → BN → ReLU → dropout → Dense(128→10).

epoch   1/8  loss 0.3010  train 91.58%  val 98.56%
epoch   8/8  loss 0.0060  train 99.91%  val 99.20%
best val 99.24% at epoch 7 -> test 99.28%  (72 errors)

CNN + augmentation — same architecture, --augment --epochs 15.

epoch   1/15  loss 0.3890  train 88.61%  val 97.88%
epoch  15/15  loss 0.0205  train 99.38%  val 99.24%
best val 99.28% at epoch 12 -> test 99.39%  (61 errors)

Training accuracy sits below validation for most of an augmented run — 88.61% against 97.88% in the first epoch above, and the MLP never closes the gap at all (97.47% train / 98.98% val at epoch 30). That inversion is the regulariser working, not a bug: the training numbers are measured on warped digits while validation sees clean ones, so they are not scored on the same task.

The MLP's worst confusion is 9 predicted as 4, eleven times — the classic MNIST failure. The conv net flattens that out: its worst is 2 predicted as 7, only five times, and no single confusion reaches double digits.

Augmentation does not spread the remaining errors evenly, it concentrates them. The augmented conv net makes 61 mistakes, and 20 of them are 9s (nine going to 4, six to 5). Its digit 3 is perfect — 1,010 test images, zero errors. What is left after augmentation is mostly the irreducibly ambiguous handwriting, which is why the worst single confusion goes up (five to nine) while the total goes down. checkpoints/cnn_aug_test_errors.png shows the mistakes it was most confident about; most are digits a human would also misread.

Runs are bit-reproducible: same --seed gives identical accuracy, because the data shuffle, weight init, and dropout masks all draw from one seeded generator.

Layout

nn/layers.py    Dense, ReLU, Tanh, Dropout, BatchNorm1d
nn/conv.py      Conv2d, MaxPool2d, Flatten, BatchNorm2d, im2col/col2im
nn/losses.py    fused softmax + cross-entropy
nn/optim.py     SGD (momentum/Nesterov), Adam, cosine schedule with warmup
nn/model.py     Sequential container, MLP/CNN factories, checkpoint save/load
data.py         MNIST download, cache, normalise, split
augment.py      affine + elastic warps, shared bilinear sampler, preview grid
train.py        training loop and CLI
evaluate.py     accuracy, per-class report, confusion matrix, error grid
gradcheck.py    numerical verification of every gradient

The layer contract

forward(x, training=True) -> out
backward(dout)            -> dx

backward follows the matching forward and reads what that pass cached. Parameter gradients are written into layer.grads and collected by Sequential.parameters() as (param, grad) pairs. Two invariants:

• Optimizers update parameters in place, so a layer's reference to its own array stays valid. Never rebind self.params[k].
• self.buffers holds state that is not learned but must survive a save/load round trip — the batch-norm running statistics. Updated in place too.

Decisions worth knowing about

Gradients are checked, not assumed. Wrong backprop is the failure mode of a from-scratch network: it still trains, just worse, and you blame the learning rate for a week. gradcheck.py compares every analytic gradient against a central finite difference in float64. All pass with ~3 orders of magnitude of margin.

What that check was missing was coverage, not tolerance. A review of this repo asked the sharper question — the checker passes, so is the math right or is the checker blind? A mutation study answered it: nine deliberately injected bugs (dropping batch norm's variance path, eps outside the sqrt, omitting the loss's 1/N, col2im overwriting instead of scatter-adding, max-pool routing to the argmin) were all caught, so the tolerance is sound. But two real holes turned up. check_end_to_end never re-seeded dropout, so a finite difference across a dropout network differenced two different networks — and every shipped config therefore pinned dropout: 0.0, meaning the architecture train.py actually trains was never end-to-end checked at all. And nothing ever called a layer with training=False, leaving the whole inference path outside the check. Both are now covered, and the shipped configs pass. The lesson is worth more than the fix: a green check is evidence about the cases it runs, and "all gradients pass" quietly meant "all gradients I thought to test".

No bias in front of batch norm. Batch norm subtracts the batch mean, which cancels a constant shift exactly — so a Dense bias before it has an identically-zero gradient and never moves. build_mlp passes bias=False there instead of carrying dead parameters. This was caught by the gradient checker, which also exposed that a plain relative error metric divides finite-difference noise by finite-difference noise when the true gradient is zero; the checker now uses an np.allclose-style absolute floor.

Softmax and cross-entropy are fused. Computed separately you take log(exp(z)) and lose precision, and the backward pass divides by a probability that may have underflowed to zero. Fused, the gradient collapses to softmax(z) - onehot(y) — no division at all.

One place divides by the batch size. Layer gradients are summed over the batch; the 1/N lives in the loss backward. Splitting that across both is how you end up with a learning rate that silently depends on batch size.

Normalisation uses training statistics everywhere. Val and test are scaled by the train mean/std, because at inference you do not get to see the test set's statistics. Borrowing them flatters the reported number.

Evaluation rebuilds the training split from the checkpoint, not from its own defaults. evaluate.py used to call data.load with its own --seed and the library default val_size, which quietly broke the discipline the previous entry describes: train with --val-size 2000, then run evaluate.py --split val, and 3000 of those 5000 "validation" images were ones the model trained on — a memorisation score reported as a validation score. A different --seed put 4590 of 5000 in the training set. The same mismatch re-standardised every split, test included, with constants the network was never trained under. The checkpoint now records the seed and val size alongside the mean and std it already stored, and evaluation reconstructs the exact split; when an older checkpoint predates those keys, the stored mean is compared against the rebuilt one and a mismatch prints a warning rather than a confident number.

Test is measured once. Model selection runs on the validation split, and the best-validation checkpoint is reloaded for a single test measurement at the end. Checking test each epoch and keeping the best gives a number you cannot reproduce.

Convolution is a matmul. A conv written as a Python loop over output pixels is unusably slow in NumPy. im2col copies every sliding window into a row of one big matrix so the whole layer becomes a single BLAS call; col2im scatters the gradients back, accumulating where windows overlap. The window extraction itself is a stride trick — as_strided reinterprets the padded buffer with a 6-D shape whose last two axes step by stride, copying nothing until the final reshape.

MaxPool2d has no padding option. im2col pads with zeros, and zero is not a neutral element for max — on an all-negative feature map the padding would win and the layer would quietly output 0 at the borders. Better to not offer the option than to offer a subtly wrong one.

Augmentation warps backwards, and never touches val or test. For each output pixel RandomAffine computes where to read from in the source and interpolates bilinearly. Warping forward instead would scatter source pixels to non-integer destinations and leave holes. Two details that are easy to get wrong: batches arrive already normalised, so pixels shifted in from outside the frame must be filled with whatever 0 maps to (-0.4241 here) rather than a literal 0.0, which would paint a grey border around every digit; and only the training split is augmented, or the reported accuracy is measured on a different distribution than the one anyone will actually use.

No flips. The augmentation that helps most on natural images poisons this dataset — a mirrored 2 is not a 2, and a flipped 6 is a 9 carrying the wrong label. python augment.py writes a preview grid so you can check the digits are still legible before training on them.

Elastic distortion is measured in pixels, not in a magic constant. The usual formulation multiplies a blurred random field by alpha ≈ 34 and pairs it with sigma = 4. Those numbers are not independent: blurring shrinks a field's variance by an amount that depends on sigma, so changing sigma silently changes the distortion strength too, and the constant means nothing on its own. ElasticDistort divides by the field's standard deviation — computed in closed form in __init__ as sqrt(1/3) · Σk², not measured per batch, so strength never wobbles with batch size — which leaves alpha as plain RMS displacement in pixels. Verified: alpha of 1, 2 and 4 measure as 0.99, 2.00 and 4.02 pixels, and unchanged across sigma.

Distortion strength was chosen by looking at it. The defaults, alpha=1.5 and sigma=6, come from a sweep, not a paper. At sigma=4 digits break down by alpha=1.5 — loops fill in and an 8 loses its waist, so the image no longer matches its label. At sigma=8 the field is smooth enough to be nearly a global shift, which RandomAffine already covers, so it adds nothing new. This is the entire reason python augment.py writes a preview grid: the first default tried here was alpha=2, sigma=4, and the grid showed it turning a 9 into a blob. Nothing in the loss curve would have said so.

Test-time augmentation averages probabilities, not logits. Averaging logits and then taking the softmax is a geometric mean of the probabilities, which lets one confidently-wrong view drag the whole ensemble; the arithmetic mean lets the other views outvote it. The clean image is always included as one of the views, because the test set is clean — averaging only over warps would discard the single view actually drawn from the test distribution. evaluate.py prints the plain and TTA accuracies together, since a TTA number with no baseline beside it cannot be checked.

Elastic distortion hurts here — in training and at test time. It is implemented, documented and off by default, because the measurement did not go the way the literature suggests it would. Over two otherwise identical 15-epoch runs it is worse than affine alone at every single epoch, by 0.2 to 1.7 validation points, ending 1.08 points down — 98.16% against 99.24% — and 98.62% against 99.39% on test, which is 138 errors against 61. Its best validation epoch is epoch 4, before the schedule has really started.

Train accuracy says why it is not simply a matter of needing longer. The elastic run ends at 98.39% train against affine's 99.38%, a gap that tracks the validation gap almost exactly rather than closing against it. The model is not overfitting and being usefully regularised — it is underfitting. The warped task is harder than this network gets around to solving, so the distortion spends capacity rather than buying generalisation.

At test time it is worse again, and more sharply. With --tta-elastic the conv net gives 50.0 errors at N=8 against 45.3 for affine views alone, and the MLP gives 108.3 — worse than not doing TTA at all, whose baseline is 107. The reason differs from the training case: every TTA view is a vote on one fixed image, so a view that lands off the data distribution is not a harder example to learn from, it is just a bad vote.

So both defaults are off, and TTA warps stay small — 5° rotation, 1 px, 5% scale, half the training amounts. --elastic and --tta-elastic stay in the CLI because a negative result you can reproduce is worth more than one you have to take on faith. What would plausibly change the verdict is a smaller alpha or a longer schedule; neither was run here, so neither is claimed.

cols.T @ d, not d.T @ cols. Same matrix mathematically, but the second form made Conv2d.backward take 1773 ms against a 17.8 ms forward. NumPy hits a pathological path when the left operand has few rows (32 output channels) against a huge reduction (~25k). Swapping the orientation took that single call from 1294 ms to 3.9 ms — 330× — and the whole training step from 1556 ms to 239 ms. The two forms agree to float32 summation-order noise (4.4e-07, against a 1.9e-05 error floor for a 25k-term float32 reduction; in float64 they agree to 6.4e-16), which was checked before trusting the swap.

Batch norm refuses a batch of one, loudly. With a single sample the batch variance is exactly zero, so the output collapses to beta, the input gradient comes out identically zero — nothing upstream learns anything — and the running variance gets dragged toward zero, wrecking every later eval-mode prediction. Every one of those is silent. It is reachable, too: the last batch of an epoch is short, and the default 55,000-example split leaves exactly one sample over at --batch-size 3, 7, 9, 21 and more. So BatchNorm raises instead, and data.batches takes a min_size that train.py sets to 2 whenever batch norm is on. PyTorch makes the same call for the same reason.

backward() follows a training forward(), and now says so. Batch norm's backward is the derivative of the batch-statistic path. After an eval-mode forward the cached statistics are the frozen running estimates, the true Jacobian is just gamma/std, and the old code applied the batch formula anyway — returning a gradient that was wrong by more than the value itself, including sign flips, with no error. The contract was already stated in a docstring; a docstring is not enforcement, so the layer now records which mode produced its cache and refuses the mismatch.

The initializer called "xavier" was LeCun. It used 1/fan_in and ignored fan_out entirely, which is LeCun init, not Xavier/Glorot's 2/(fan_in + fan_out). For Conv2d(1, 16, 5) that is 2.9x the Glorot scale. Rather than silently change every trained model, lecun is now its own name and the builders ask for it explicitly — so the numbers in this README are the numbers you reproduce — while xavier now means actual Glorot for anyone who asks for it. Conv2d also used to accept any unrecognised init string and quietly fall through to the LeCun branch, where Dense raised; both raise now.

Component speedups did not compound into the end-to-end win they promised. A profiling pass found real waste and the fixes all landed: batch norm's backward was recomputing two full-tensor reductions that are algebraically gamma·dbeta/n and gamma·dgamma/n, quantities the two lines above it had just computed; BatchNorm2d transposed NCHW into a 2-D view and back on both passes, immediately after Conv2d had transposed into NCHW; max pooling pushed non-overlapping windows through im2col when a plain reshape gives the same windows as free views; and Adam allocated about six full-size temporaries per parameter per step where two folded scalars and one scratch buffer do. Each was measured, and each is bit-identical or float32-noise-identical to what it replaced.

Added up as isolated micro-benchmarks they predicted roughly a third off the training step. The actual, interleaved, paired measurement is 1.09x median — faster in 7 of 8 rounds, and once slower. Two reasons, both worth knowing. Micro-benchmarks of allocation-heavy code flatter the rewrite, because a tight loop reuses a warm allocator that a real training step keeps cold. And Amdahl: what remains is the conv matmul and col2im, which together dominate the step and were already about as good as pure NumPy gets. The honest summary is that this bought about 9% of wall time and a much tidier BatchNorm — the win worth having was elsewhere, in memory.

Inference was holding half a gigabyte of dead activations. Every layer cached what its backward pass needs — Conv2d its im2col matrix, batch norm its normalised input — with no regard for whether a backward pass was ever coming. After one validation pass at the default eval batch of 1000, the model was pinning 493 MB of caches that nothing would ever read, and holding them through the whole of the next training epoch. Gating each cache on training takes that to zero, which is the single largest effect of the whole performance pass and cost one conditional per layer.

A separable blur wider than the image is not separable, it is dense. The elastic field is smoothed with a Gaussian truncated at 4σ, so sigma=6 gives a 49-tap kernel — on a 28-pixel image. Sliding that kernel means 2 × 49 passes over the whole batch, and every pass past the 28th is multiplying by zeros outside the frame. Since the blur is a fixed linear operator, it can be built once as a 28 × 28 matrix per axis and applied as (By @ f) @ Bx.T: identical arithmetic, handed to BLAS instead of a Python loop. That is 15× on the blur, and it takes affine+elastic from 110 s to 30 s per epoch — the difference between a four-hour run and a half-hour one. The two forms agree to 3.4e-07 relative, which is float32 summation-order noise, and alpha=1.5 still measures 1.50 px of RMS displacement at the default sigma, so the rewrite changed the cost and nothing else. Worth noticing early: the shape of the data decides whether a "separable" optimisation is one at all.

Options

--arch mlp|cnn         --epochs (20 mlp / 8 cnn)   --batch-size 128
--lr 1e-3              --optimizer adam|sgd        --momentum 0.9
--weight-decay 0.0     --activation relu|tanh      --dropout 0.2
--no-batchnorm         --warmup 100                --seed 0
--out PATH             --plot

mlp:  --hidden 256 128
cnn:  --channels 16 32   --kernel-size 5   --fc-hidden 128
aug:  --augment   --aug-rotation 10 (deg)   --aug-translate 2 (px)
                  --aug-scale 0.1 (fraction)
      --elastic   --elastic-alpha 1.5 (px RMS)   --elastic-sigma 6 (px)

--augment and --elastic are independent and compose; either can be used alone. evaluate.py takes its own set:

--tta N   --tta-rotation 5   --tta-translate 1   --tta-scale 0.05
          --tta-elastic

Augmentation makes the training task harder, so it wants more epochs than the defaults: the results above use 30 for the MLP and 15 for the CNN, against defaults of 20 and 8. Expect train accuracy to sit below validation accuracy while it is on — that is the regulariser working, not a bug.

--out defaults to checkpoints/<arch>.npz, and evaluate.py with no --checkpoint picks whichever was trained most recently. Checkpoints store their own architecture config, so evaluation rebuilds the right model — and picks the right data layout, flat or NCHW — without being told.

Some things to try:

python mnist_nn_scratch/train.py --optimizer sgd --lr 0.1 --no-batchnorm

python mnist_nn_scratch/train.py --arch cnn --channels 32 64 --kernel-size 3 --epochs 15

Extending

To add a layer, subclass Layer, implement forward/backward, register arrays in self.params/self.grads, then add a case to gradcheck.py — the checker is the thing that tells you whether the new backward pass is right. Every conv layer here passed on the first run because the checker existed before the layers did.

Natural next steps: average pooling and global average pooling (easy, and col2im already does the scatter); residual connections (needs a container that is not a straight line, so Sequential would grow a sibling); and an ensemble of independently seeded models, which is the obvious thing left after TTA — TTA averages one model over several views of an image, and an ensemble averages several models over one view.