MNIST Transformer in PyTorch (Muon + WSD)

A PyTorch port of the best training configuration from david-yu/neural-network — a transformer built from scratch in Rust with manual backpropagation. This repo reimplements that project's muon_wsd_b8 recipe on autograd to (1) verify the hand-derived gradients against a reference implementation and (2) search the recipe space cheaply. Everything is one file: muon_wsd_b8.py.

What it trains

The exact model and data pipeline of the Rust repo:

• Model: encoder-only transformer classifier with Full Attention Residuals (AttnRes, Kimi Team) in place of standard residual connections — every sub-layer and the output head read a learned softmax mixture over the token embedding and all earlier sub-layer outputs. 2 layers, d_model 64, 4 heads, d_ff 128, sinusoidal positions, mean pooling: 69,066 parameters (asserted at startup).
• Data: MNIST as 49 tokens — a 7x7 grid of 4x4 average-pooled patches quantized to 16 gray levels, parsed from the raw IDX .gz files. Tokenization is integer-exact and byte-identical to the Rust implementation (verified sample-for-sample). PATCH and LEVELS relax this to finer grids; see the tokenizer section below for what each one buys and costs.
• Recipe (muon_wsd_b8): Muon (Newton-Schulz orthogonalized Nesterov momentum, Moonlight's 0.2*sqrt(max(m,n)) RMS matching) on the six hidden matrices per layer, Adam on everything else, the AttnRes pseudo-queries at 0.1x, a Warmup-Stable-Decay schedule, batch 8, lr 0.001*sqrt(8), 10 epochs on 20,000 train / 2,000 test samples.

Parity with the Rust implementation

Same recipe, 10 epochs, 3 runs each (Rust runs are its own RNG; PyTorch runs are seeds 0/1/2 — comparable in distribution, not paired):

Implementation	Final test acc	Runs
Rust (manual backprop)	94.45%	95.05 / 93.80 / 94.50
PyTorch (this repo)	94.05%	93.90 / 94.45 / 93.80

Overlapping spreads around ~94: the hand-derived gradients and autograd agree. A batch-8 epoch takes ~12s on an Apple M5 CPU.

Recipe search: what actually improves it

All runs 20,000 train / 2,000 test, batch 8 unless noted, 2-3 seeds each.

Change vs parity	Final test acc	Test loss	Verdict
— (parity, 10 epochs)	94.05%	0.197	baseline
30 epochs	94.12%	0.303	null — overfits: 3x the training, same accuracy, worse loss
30 epochs, batch 32 (lr x sqrt(batch))	94.00%	0.278	null
30 epochs, Muon weight decay 0.1	94.40%	0.176	small win — fixes the overfitting (Moonlight's finding reproduced)
OneCycle schedule, 2x peak lr	94.03%	0.198	null
OneCycle schedule, 4x peak lr	91.50%	0.266	negative — the hot-peak trick from CNN speedruns does not transfer
Shift augmentation ±2px, 10 epochs	95.47%	0.145	+1.4 — seeds within ±0.05 of each other
Shift augmentation ±2px, 30 epochs	95.75%	0.139	+1.7 — augmentation converts the extra epochs from overfitting into gains
Shift augmentation ±2px + wd 0.1, 30 epochs	95.22%	0.155	negative vs augmentation alone — the two regularizers overlap
Shift augmentation ±2px, 60 epochs	96.08%	0.124	best: +2.0 over parity (per-epoch bests reached 96.60)

Best measured configuration (+2.0 points over the Rust recipe, ~15 min on an M-series CPU):

EPOCHS=60 AUGMENT=2 python muon_wsd_b8.py

The trajectory 94.05 -> 95.47 -> 95.75 -> 96.08 over 10/30/60 epochs shows what the augmentation actually changed: without it, training past 10 epochs only overfits; with it, more compute keeps converting into accuracy (with diminishing returns by 60).

The augmentation is random integer pixel shifts of the 28x28 image before tokenization (the transferable core of the 99.7%-class Kaggle MNIST recipes, which shift/rotate/zoom) — the model, tokenizer and evaluation stay identical while effective training data multiplies.

Where the remaining gap to ~99% lives — measured. That the tokenizer is what caps this task used to be an assumption here. It is now a measurement. Holding one CNN and one training recipe fixed and changing only what the model is allowed to see (60,000 train / the full 10,000 test images, 30 epochs):

Input representation	Test acc
28x28, 256 gray levels (raw pixels)	99.67%
14x14, 16 levels	99.52%
7x7, 16 levels — this repo's tokenizer	98.15%
7x7, 256 levels	98.27%
7x7, 16 levels, + TTA + 3-model ensemble	98.84%

Three things follow. 99.5% is unreachable through 49 tokens — the ceiling is ~98.8%, and that is with a 1.68M-parameter CNN, 60k images and every trick applied. The cost is spatial resolution, not quantization: 7x7 at full 256-level depth buys only +0.1 over 16 levels, while 14x14 at the same 16 levels clears 99.5%. So the 16-bin quantization is nearly free and the 4x4 pooling is what discards the digits. And the tokenizer was not the binding constraint at 96.08% — the representation supports 98.15% on its own, so about two points were sitting in training rather than in the input spec. representation_ceiling.py reproduces the table.

Tokenizer resolution: what a finer grid buys, and what it costs

PATCH makes the tokenizer a dial rather than a constant. PATCH=4 is the 49-token parity tokenizer, PATCH=2 gives 196 tokens, PATCH=1 gives the raw 784 pixels (pair it with LEVELS=256 to keep every gray value). Same model, same recipe, same data — 20,000 train / the full 10,000 test, 20 epochs, batch 8, shift augmentation, 2 seeds each:

Tokenizer	Tokens	Test acc	Wall clock / run
7x7 (PATCH=4, parity)	49	96.91%	12 min
14x14 (PATCH=2)	196	98.41%	29 min
28x28 (PATCH=1 LEVELS=256)	784	98.06%	91 min

PATCH=2 is the best-returning change measured in this repo: +1.5 points for 2.4x the compute. It also beats the 49-token ceiling of 98.15% — the 69k transformer reading 196 tokens outperforms the strongest model measured on 49 tokens, a gap no amount of training could have closed.

Raw pixels are a regression, not a win. PATCH=1 hands the model strictly more information than PATCH=2 and strictly more compute — 7.4x parity, 3.2x PATCH=2 — and finishes 0.35 points behind it. Comparing each tokenizer against what its own representation supports shows the model, not the input, giving way:

Tokenizer	This model	Representation ceiling	Gap
7x7, 49 tokens	96.91%	98.15%	−1.24
14x14, 196 tokens	98.41%	99.52%	−1.11
28x28, 784 tokens	98.06%	99.67%	−1.61

The gap to the ceiling widens at 784 tokens: the extra resolution is there and the model cannot use it. Two reasons, both structural. The positional encoding is 1-D sinusoidal over a row-major flattening, so nothing tells the model that pixel i and pixel i+28 are vertically adjacent — a mild handicap over 49 patches, a severe one over 784 pixels. And 2 layers with mean pooling have to summarize 16x more tokens on the same 20-epoch budget. Fixing this means changing the architecture (2-D positions, more depth), not the tokenizer, so PATCH=2 is where the returns are until then.

Why the larger grids cost what they do

Attention is O(T^2) while every other part of the model is O(T), so cost grows faster than the token count and the mix shifts as the grid gets finer (per layer, per sample, d_model 64 / d_ff 128):

Tokens	Attention share of a layer	Fwd FLOPs / sample	Scores tensor at batch 8
49	16%	3.8M	0.3 MB
196	43%	22.7M	4.9 MB
784	75%	209M	79 MB

Measured wall clock is kinder than that 54x FLOP ratio — one epoch over 2,000 samples at batch 8 takes 2s / 5s / 24s on CPU for 49 / 196 / 784 tokens, so 2.5x and 12x rather than 5.9x and 54x. At 49 tokens the matmuls are small enough to be latency-bound, so the larger grids use the hardware better than their FLOP count suggests. Two practical consequences at PATCH=1: evaluation batches shrink automatically (1000 -> 13) to hold the 784x784 scores tensor to a fixed memory budget, which makes each eval pass 769 forward calls; and DEVICE=mps finally wins over CPU (16s vs 24s per epoch), where at 49 tokens it loses badly (15s vs 2s) because the model is too small to cover kernel-launch overhead.

Note that LEVELS changes the embedding table, and with it the parameter count: LEVELS=256 makes the model 84,426 parameters instead of 69,066, so the parity assertion applies only to the default PATCH=4 LEVELS=16 configuration.

Usage

bash download_mnist.sh          # fetches the four MNIST .gz files into data/
pip install torch
python muon_wsd_b8.py           # parity configuration

Configuration via environment variables:

Variable	Default	Meaning
TRAIN, TEST	20000, 2000	dataset subset sizes
EPOCHS	10	training epochs
BATCH	8	batch size (lr scales by sqrt(batch))
LR	0.001	base learning rate, pre-scaling
SCHEDULE	wsd	wsd or onecycle
AUGMENT	0	max random pixel shift; 0 disables
PATCH	4	tokenizer pooling window: 4 -> 49 tokens, 2 -> 196, 1 -> 784
LEVELS	16	gray levels per token (use 256 with PATCH=1 for raw pixels)
WD	0.0	Muon decoupled weight decay
SEEDS	0,1,2	comma-separated run seeds
DEVICE	cpu	cpu or mps

License

GPL-3.0, matching the parent repository.