MLP on MNIST: A Complete Local-Only NumPy Recipe

0 · MNIST as an instrument, not a promised score

MNIST contains 60,000 training and10,000 test 28×28 grayscale handwritten digits. A small MLP is an inspectable way to connect data, initialization, stable loss, backpropagation and evaluation. A 97% score or a90-second runtime is not an invariant of the architecture. This review did not download or train on MNIST. The complete local-only recipe below consumes files you already possess; the executed demonstration uses explicitly synthetic data.

The original cheque-reader story mixed a named commercial system, throughput and market-share figures without support. LeCun and colleagues' document-recognition work provides historical context, not a claim that this MLP matches LeNet-5 or that 99.7% is an accuracy ceiling. Likewise, grokking on modular arithmetic is not an MNIST experiment. A small research probe is useful because its variables can be controlled, not because passing one benchmark proves a new idea sound.

1 · What you'll build today

Read local IDX files, inspect shapes/pixels/labels, derive a 784→128→10 MLP, implement stable softmax cross entropy and every gradient, shuffle paired examples into minibatches, select a checkpoint using validation only, and reserve the test set for final reporting. You will also design seven ablations and distinguish the actual tasks in grokking, lottery-ticket and connectivity papers. Correctness takes priority over squeezing the script into48 lines.

2 · Checkpoint · you should already know

Review probability and information, optimization, forward shapes and backprop. Optional visual introductions are 3Blue1Brown, Samson Zhang, and Nielsen. The MNIST reference is historical background; the verified dataset-format evidence is the tagged torchvision reader, not a claimed successful download from that site.

3 · Load MNIST — local files, no framework or network

Place already-owned raw or gzip IDX files in a directory: train-images-idx3-ubyte, train-labels-idx1-ubyte, t10k-images-idx3-ubyte, t10k-labels-idx1-ubyte (optionally .gz). The image magic is2051 and label magic2049; headers use big-endian32-bit integers. Validate exact payload size before reshape, image geometry, matching counts and label range. Missing files raise rather than trigger a download. Do not substitute the scikit-learn8×8 digits dataset and label its accuracy “MNIST.”

import gzip
import struct
from pathlib import Path
import numpy as np
 
def read_idx(path, image):
    path = Path(path)
    raw = gzip.open(path,'rb').read() if path.suffix=='.gz' else path.read_bytes()
    header = 16 if image else 8
    if len(raw)<header:
        raise ValueError('short IDX header')
    fields=struct.unpack('>'+'I'*(header//4),raw[:header])
    magic,n=fields[:2]
    if magic!=(2051 if image else 2049) or n==0:
        raise ValueError('IDX magic/count')
    if image and fields[2:]!=(28,28):
        raise ValueError('expected 28 by 28')
    expected=n*(784 if image else 1)
    if len(raw)-header!=expected:
        raise ValueError('IDX payload length')
    a=np.frombuffer(raw,dtype=np.uint8,offset=header).copy()
    return a.reshape(n,784) if image else a
 
def load_local_mnist(root):
    root=Path(root)
    def find(name):
        p=root/name
        return p if p.exists() else root/(name+'.gz')
    arrays=[]
    for split,expected in [('train',60000),('t10k',10000)]:
        X=read_idx(find(split+'-images-idx3-ubyte'),True)
        y=read_idx(find(split+'-labels-idx1-ubyte'),False)
        if len(X)!=len(y) or len(X)!=expected or np.any(y>9):
            raise ValueError('MNIST count/label mismatch')
        arrays.extend([X.astype(np.float32)/255.,y.astype(np.int64)])
    return tuple(arrays)

3.1 Visualise a few — always look at your data

A dependency-free grayscale SVG makes eight images inspectable without installing a plotting library. Save only to your chosen output location. In the audited synthetic run this displays synthetic pixels, not downloaded digits. For real local MNIST, inspect ambiguous digits and preprocessing errors without asserting an unknowable label-noise ceiling.

def image_grid_svg(X,y):
    out=['<svg xmlns="http://www.w3.org/2000/svg" width="896" height="132" viewBox="0 0 224 33">']
    for j,row in enumerate(X[:8]):
        pix=np.clip(np.rint(row.reshape(28,28)*255),0,255).astype(int)
        out.append(f'<text x="{28*j}" y="4" font-size="4">{int(y[j])}</text>')
        for r in range(28):
            for c in range(28):
                q=int(pix[r,c]);out.append(f'<rect x="{28*j+c}" y="{r+5}" width="1" height="1" fill="rgb({q},{q},{q})"/>')
    return ''.join(out)+ '</svg>'

3.2 Sanity checks and memory

Check shape, dtype, finite values, min/max, mean/std and np.bincount(y,minlength=10) from the actual arrays. The70,000×784 float32 pixels alone occupy 219,520,000bytes ≈209.35MiB, not180MB; labels, original bytes, shuffled copies and activations add memory. The original's stated count vector [5923,6742,5958,6131,5842,5421,5918,6265,5851,5949] sums to60,000; max/min≈1.244, so “within10% across classes” is not an accurate max/min claim. Compute actual class counts from your files; this vector is not a measurement from the audit. Neither mean intensity.13 nor class balance tells you what fraction of pixels is exactly zero. A uniform-random classifier has expected accuracy10%; a majority-class baseline depends on measured frequencies.

4 · The 784-128-10 MLP with softmax

Z1=X@W1+b1, H=ReLU(Z1), Z2=H@W2+b2, P=softmax(Z2). For B rows the shapes are (B,784)→(B,128)→(B,10). The parameters number 784*128+128+128*10+10=101770. Batch size changes activation storage and gradient estimates, not the number of parameters.

4.1 He initialization — second moments, not halved variance

For independent zero-mean weights and approximately identically distributed inputs, E[z²]=fan_in*Var(W)*E[x²]. A symmetric preactivation distribution gives E[ReLU(z)²]=E[z²]/2. Preserve this second moment by choosing Var(W)=2/fan_in. ReLU's mean is positive, so its variance is not simply half the preactivation variance. For Gaussian z with variance q, Var(ReLU(z))=q*(1/2−1/(2π)). Assumptions can degrade during training.

The final linear layer need not use the ReLU gain. We use 1/fan_in there and expose sizes/seed rather than relying on global NumPy RNG. Xavier's general fan-in/fan-out rule is not synonymous with 1/fan_in in every shape.

def init_params(sizes=(784,128,10),seed=0):
    rng=np.random.default_rng(seed)
    d,h,k=sizes
    return (rng.normal(size=(d,h))*np.sqrt(2/d),np.zeros(h),
            rng.normal(size=(h,k))/np.sqrt(h),np.zeros(k))

5 · Forward pass and stable log probabilities

Subtracting a row maximum leaves softmax unchanged because numerator and denominator share the same exponential factor. It prevents overflow for finite logits, but tiny probabilities can still underflow. Compute CE from logits rather than log(P+epsilon). For [1000,-1000] and target1, the correct loss is2000; an epsilon-capped probability loss incorrectly reports a small fixed ceiling.

def softmax_stable(z):
    shifted=z-z.max(axis=1,keepdims=True)
    e=np.exp(shifted)
    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_stable(Z2)
    return P,(X,Z1,H,Z2,P)
 
def cross_entropy_logits(z,y):
    if len(y)==0 or z.ndim!=2 or y.shape!=(len(z),):
        raise ValueError('nonempty (batch,classes) logits and flat labels required')
    if not np.issubdtype(y.dtype,np.integer) or np.any(y<0) or np.any(y>=z.shape[1]):
        raise ValueError('class indices out of range')
    shifted=z-z.max(1,keepdims=True)
    return float(np.mean(np.log(np.exp(shifted).sum(1))-shifted[np.arange(len(y)),y]))
assert cross_entropy_logits(np.array([[1000.,-1000.]]),np.array([1]))==2000
assert np.isclose(cross_entropy_logits(np.zeros((3,10)),np.array([1,2,3])),np.log(10))

6 · Loss and backward — why P minus one-hot appears

For one example, L=log(sum(exp(z)))-z[y]. Differentiate each logit to get P[j]-1[j=y]; the mean contributes1/B. Mutate a copy of P so the cached probabilities remain available. Avoiding an explicit one-hot allocation saves one B×K array, not10× of the whole layer's memory.

def backward(y,cache,params):
    X,Z1,H,Z2,P=cache
    W1,b1,W2,b2=params
    N=y.size
    if N==0 or y.shape!=(len(X),):
        raise ValueError('nonempty aligned labels required')
    dZ2=P.copy()
    dZ2[np.arange(N),y]-=1
    dZ2/=N
    dW2=H.T@dZ2;db2=dZ2.sum(axis=0)
    dH=dZ2@W2.T;dZ1=dH*(Z1>0)
    dW1=X.T@dZ1;db1=dZ1.sum(axis=0)
    return dW1,db1,dW2,db2

This resembles sigmoid/BCE's prediction-minus-target cancellation. It does not cover every classification objective unchanged: class weights, focal factors, label smoothing, masks and reduction semantics require their own derivations.

7 · Training loop — paired shuffle and validation selection

A permutation keeps X and y paired. Sample minibatches by indices instead of copying the whole dataset. Retain the final partial batch and divide by its actual size. For60,000 examples at64, there are ceil(60000/64)=938 updates per epoch, thus9,380 over10 epochs; a55,000-example training subset has860 per epoch.

def metrics(X,y,p,batch_size=512):
    total,correct=0.,0
    if len(y)==0: raise ValueError('empty evaluation')
    for s in range(0,len(y),batch_size):
        xb,yb=X[s:s+batch_size],y[s:s+batch_size]
        P,cache=forward(xb,p)
        total+=len(yb)*cross_entropy_logits(cache[3],yb)
        correct+=int((P.argmax(1)==yb).sum())
    return total/len(y),correct/len(y)
 
def train(X,y,X_val,y_val,epochs=10,batch_size=64,lr=.1,seed=0,hidden=128):
    if batch_size<=0 or len(y)==0: raise ValueError('empty data or invalid batch')
    rng=np.random.default_rng(seed)
    p=init_params((X.shape[1],hidden,int(max(y.max(),y_val.max()))+1),seed)
    best=None;best_loss=float('inf');history=[]
    for epoch in range(epochs):
        idx=rng.permutation(len(y))
        for s in range(0,len(y),batch_size):
            j=idx[s:s+batch_size];_,cache=forward(X[j],p)
            g=backward(y[j],cache,p)
            p=tuple(a-lr*b for a,b in zip(p,g))
        tr=metrics(X,y,p);va=metrics(X_val,y_val,p)
        if not np.isfinite(tr[0]+va[0]): raise FloatingPointError('nonfinite loss')
        history.append((tr[0],va[0],va[1]))
        if va[0]<best_loss:
            best_loss=va[0];best=tuple(a.copy() for a in p)
    if best is None: raise ValueError('positive epochs required')
    return best,np.asarray(history)

For full MNIST assert all10 classes occur in the training split; in a general classifier pass an explicit class count instead of inferring unseen classes from validation labels. Here class IDs are dataset metadata, never fitted from test labels. A validation curve is for model selection; do not inspect test accuracy after every ablation and still call the test set untouched.

8 · Read and plot the loss curve

Uniform logits give ln(K) independent of label balance. Random weights do not guarantee uniform logits, so a different step-zero loss alone is not a bug. A plateau may be optimization failure, capacity limit or noise, not proven convergence. Spikes suggest checking learning rate, outliers, data order and numeric range; NaN can originate before softmax. Decreasing training loss with rising validation loss is evidence of overfitting under the fixed evaluation protocol, not proof that dropout is the right fix.

def loss_curve_svg(history):
    ymax=max(float(history[:,:2].max()),1e-12)
    out=['<svg xmlns="http://www.w3.org/2000/svg" width="480" height="180" viewBox="0 0 480 180">']
    for j,color in [(0,'blue'),(1,'orange')]:
        pts=' '.join(f'{10+460*i/max(1,len(history)-1):.2f},{170-150*v/ymax:.2f}' for i,v in enumerate(history[:,j]))
        out.append(f'<polyline points="{pts}" fill="none" stroke="{color}"/>')
    return ''.join(out)+'</svg>'

Blue is training loss; orange is validation. Save the epoch values alongside the plot so visual smoothness does not replace numbers. Accuracy uses counts across examples, and mean loss weights each batch by its actual size.

9 · Seven ablations — hypotheses, not fabricated outcomes

Hold data split and initialization policy fixed; repeat seeds and log samples processed, updates, wall time and validation loss. Compare one change at a time:

1. Rate .1→1: test overshoot versus faster progress; divergence is possible, not mandatory.
2. Batch1: more, noisier updates and greater interpreter/launch overhead; no guaranteed generalization improvement.
3. Batch1024: fewer updates per epoch and larger matrices; measure throughput separately from convergence. Linear learning-rate scaling is a heuristic under approximately unchanged gradients, not a law.
4. Tiny initialization: inspect activation/gradient second moments; do not promise a2–3× slowdown or fixed accuracy penalty.
5. Identity activation: collapse to a linear-logit classifier; use its measured validation score as a baseline rather than assuming92%.
6. Additional layer or width: account for extra parameters/FLOPs and compare matched budgets. No guaranteed half-percent gain.
7. No shuffle: measure order sensitivity and inspect actual labels. MNIST is not generally stored as all zeroes then all ones, and SGD need not shuffle under every valid sampling scheme.

For independent per-sample gradients, the mean has variance proportional to1/B, so standard deviation scales as1/sqrt(B); finite-population sampling has a correction. This is a variance statement at a fixed parameter point, not a theorem that small batches generalize better. Matching many small updates to one big update assumes gradients change little. On a positive-definite quadratic the stability bound is 0<lr<2/lambda_max; a changing nonlinear Hessian is not one global safe rate. Goyal et al.'s large-batch ImageNet result used a particular warmup/training regime, not this exact MNIST program.

Choose larger batches for measured throughput or communication constraints, higher learning rate for validated optimization headroom, or input-pipeline work for measured stalls. Low utilization alone does not identify which one. Plotting loss against epochs, updates, samples and time answers different questions; none should be silently substituted for another.

10 · The complete local-only recipe and synthetic contract

Concatenate the Python fences in order. Then call run_local_mnist yourself only if the files already exist. It produces a best-validation checkpoint, one final test evaluation, pixel and loss SVGs, and a confusion matrix. The test set is not used to pick epochs. Normalizing by255 is a fixed physical pixel-scale conversion; any learned mean/std must be fitted on training only.

def run_local_mnist(root,out_dir):
    X,y,Xtest,ytest=load_local_mnist(root)
    rng=np.random.default_rng(19);idx=rng.permutation(len(y))
    ti,vi=idx[:55000],idx[55000:]
    assert np.unique(y[ti]).size==10
    best,history=train(X[ti],y[ti],X[vi],y[vi])
    test_loss,test_acc=metrics(Xtest,ytest,best)
    pred=forward(Xtest,best)[0].argmax(1)
    confusion=np.zeros((10,10),dtype=np.int64)
    np.add.at(confusion,(ytest,pred),1)
    out=Path(out_dir);out.mkdir(parents=True,exist_ok=True)
    np.savez(out/'mnist-mlp.npz',**dict(zip(('W1','b1','W2','b2'),best)))
    np.savetxt(out/'history.csv',history,delimiter=',',header='train_loss,val_loss,val_accuracy')
    np.savetxt(out/'confusion.csv',confusion,delimiter=',',fmt='%d')
    (out/'pixels.svg').write_text(image_grid_svg(X[ti],y[ti]))
    (out/'loss.svg').write_text(loss_curve_svg(history))
    print('LOCAL MNIST TEST',test_loss,test_acc,'examples',len(ytest))
    return best,history,confusion
 
# Executed locally: synthetic separable classes, NOT MNIST training.
rng=np.random.default_rng(7)
sy=np.repeat(np.arange(3),20)
sx=np.eye(3)[sy]+rng.normal(0,.08,(60,3))
vp,vh=train(sx[::2],sy[::2],sx[1::2],sy[1::2],epochs=30,batch_size=7,lr=.2,hidden=8)
sl,sa=metrics(sx[1::2],sy[1::2],vp)
assert sa>.95 and vh[-1,0]<vh[0,0]
assert '<polyline' in loss_curve_svg(vh)
print('SYNTHETIC ONLY',sl,sa,'epochs',len(vh))

The synthetic fixture exercises training, uneven batches, selection and evaluation. Separate tests exercise exact local IDX parsing with tiny generated records and malformed-header controls. They do not claim the60k/10k training result. Training computations default to float64 parameters in this teaching code despite float32 image storage; record that dtype choice in timing/memory reports.

11 · Small-scale research: three distinct experiments

Grokking: Power et al. study small algorithmic datasets; Nanda et al. analyze a one-layer transformer learning modular addition through Fourier features. Delayed test improvement after memorization motivates tracking train/test behavior over time, but neither makes the present MNIST MLP their principal experiment.

Lottery tickets: Frankle and Carbin identify pruned, initialized subnetworks that train effectively under their protocol. The initialization and pruning/reset procedure matter; “any sparse mask works” is not the claim. MNIST is a real part of that experimental lineage.

Linear mode connectivity: Frankle et al. use interpolation loss barriers to study stability under SGD noise and pruning, including early-training behavior. Same initialization alone does not guarantee that arbitrary endpoints lie in a flat basin. To test two compatible parameter sets, evaluate theta(alpha)=(1-alpha)*thetaA+alpha*thetaB on a fixed held-out split and compare the curve to endpoint loss; align architecture and parameter coordinates first. Model soups are a related averaging method, not a universal consequence of a toy interpolation plot.

12 · Failure cases

Unnormalized pixels amplify preactivations and gradients; divide uint8 intensities by255 here, but values above1 are not universally invalid features. Shifted softmax prevents exponential overflow, not a nonfinite matmul. Swapping train/test names corrupts the protocol but does not mathematically force12% accuracy. fetch_openml does not inject an undocumented training RNG into this local recipe. Use explicit generators and record library versions, array hashes and BLAS/thread settings; Python dictionaries preserve insertion order, so a hash seed is not a generic fix for parameter initialization.

Pair shuffle indices with labels; compare per-class counts and sampled IDs. If a run is unusually fast, verify the number of records and updates. Distinguish data bugs from optimization failure before tuning the architecture.

13 · Try it yourself — answers and acceptance criteria

• Validation split: §10 uses55k/5k and copies best arrays, avoiding a live reference that changes during later updates. Plot both curves; an overfitting turn is possible, not mandatory.
• Confusion: rows are true labels, columns predicted labels. Remove the diagonal before finding the largest error pair; inspect actual pixels instead of presuming4↔9 or3↔5.
• Freeze W1: after epoch5 update only indices2,3 of the parameter tuple. Features H become fixed; the last layer learns linear logits in that feature space. Accuracy need not improve, and this is not full fine-tuning.
• Depth/width: replace the forward/backward with a list of affine caches and reverse traversal, count the extra parameters, and select on validation. “Beat98% in two lines” omitted the required backward changes.
• Fashion-MNIST: use already-local files with the same IDX geometry and its clothing-label meanings; create a separate run directory and split manifest. Do not mix digit labels, assert a10-point accuracy gap or claim a convolution will universally win without a controlled comparison.

A worked ablation engine — synthetic data, all seven switches

Here is the missing implementation behind the ablation instructions. It supports arbitrary hidden widths, identity versus ReLU, tiny initialization, batch size/rate, shuffling and freezing all hidden layers after a chosen epoch. All variants below use the same synthetic split from §10. Their printed validation values are not MNIST or Fashion-MNIST accuracies and do not establish a universal ranking.

def general_eval(X,y,layers,activation='relu'):
    caches=[];a=X
    for i,(W,b) in enumerate(layers):
        z=a@W+b;caches.append((a,z))
        a=np.maximum(0,z) if i<len(layers)-1 and activation=='relu' else z
    P=softmax_stable(a);dz=P.copy();dz[np.arange(len(y)),y]-=1;dz/=len(y)
    grads=[]
    for i in range(len(layers)-1,-1,-1):
        inp,z=caches[i];W,b=layers[i]
        grads.append((inp.T@dz,dz.sum(0)))
        if i:
            dz=dz@W.T
            if activation=='relu':dz*=caches[i-1][1]>0
    return cross_entropy_logits(a,y),list(reversed(grads)),P
 
def ablate(X,y,Xv,yv,hidden=(8,),batch_size=7,lr=.2,epochs=20,
           activation='relu',tiny=False,shuffle=True,freeze_after=None):
    rng=np.random.default_rng(7);widths=[X.shape[1]]+list(hidden)+[3]
    layers=[(rng.normal(size=(a,b))*(.01 if tiny else np.sqrt(2/a)),np.zeros(b))
            for a,b in zip(widths,widths[1:])]
    snapshot=None;steps=0
    for epoch in range(epochs):
        if freeze_after is not None and epoch==freeze_after:
            snapshot=[(W.copy(),b.copy()) for W,b in layers[:-1]]
        ids=rng.permutation(len(y)) if shuffle else np.arange(len(y))
        for start in range(0,len(y),batch_size):
            j=ids[start:start+batch_size];_,g,_=general_eval(X[j],y[j],layers,activation)
            for i,((W,b),(dW,db)) in enumerate(zip(layers,g)):
                if freeze_after is None or epoch<freeze_after or i==len(layers)-1:
                    W-=lr*dW;b-=lr*db
            steps+=1
    if snapshot is not None:
        assert all(np.array_equal(W,A) and np.array_equal(b,B) for (W,b),(A,B) in zip(layers[:-1],snapshot))
    value,_,P=general_eval(Xv,yv,layers,activation)
    assert np.isfinite(value)
    return {'val_loss':value,'val_accuracy':float((P.argmax(1)==yv).mean()),
            'updates':steps,'parameters':sum(W.size+b.size for W,b in layers)}
 
variants={'baseline':{},'rate1':{'lr':1.},'batch1':{'batch_size':1},
          'batch1024':{'batch_size':1024},'tiny_init':{'tiny':True},
          'identity':{'activation':'identity'},'extra_layer':{'hidden':(8,6)},
          'no_shuffle':{'shuffle':False},'freeze_hidden':{'freeze_after':5}}
results={name:ablate(sx[::2],sy[::2],sx[1::2],sy[1::2],**opts) for name,opts in variants.items()}
assert results['baseline']['updates']==20*5
assert results['batch1']['updates']==20*30
assert results['batch1024']['updates']==20
assert results['extra_layer']['parameters']>results['baseline']['parameters']
print('SYNTHETIC ABLATIONS ONLY',results)

The engine fixes three classes because this fixture's class vocabulary is known; for MNIST change that configured endpoint to10. A frozen hidden-layer snapshot is asserted unchanged, making the linear-probe exercise testable. One20-epoch comparison is not a hyperparameter study: repeat seeds and compare wall time/samples as described in §9 before drawing an empirical conclusion.

14 · Recap and retention answers

1. Stable softmax: subtract each row maximum, exponentiate, divide by row sums; do not apply a global maximum or wrong axis.
2. CE backward: copy P, subtract1 at [arange(N),y], divide by N. The copy preserves the forward cache.
3. He factor: symmetric preactivation halves the second moment after ReLU; compensate with weight variance2/fan_in, not standard deviation2/fan_in.
4. **Count:**101,770 trainable scalars.
5. Shuffle: paired random permutations reduce persistent order correlations; no predetermined unshuffled accuracy and no claim of label-sorted MNIST.
6. Uniform loss: -log(1/K)=log(K); random logits need not be uniform.
7. Research: the2020 connectivity paper investigates SGD-noise/pruning stability using interpolation barriers; Nanda's2023 probe is a small modular-addition transformer, not MNIST. Small scale permits controlled repetition, not proof that larger-scale study is impossible.
8. Fashion result: no score from this review. The procedure is local loading, separate validation selection, then one final test report. Compare architecture changes under a stated budget rather than assuming a winner.

Reimplement the stable endpoint tomorrow; rerun a controlled synthetic ablation next week. The MNIST experiment remains ready for local execution, not falsely reported as completed.

Sources and execution boundary

Primary references: torchvision v0.24 MNIST/IDX reader, He et al. 2015, PyTorch2.9 cross entropy, NumPy2.2 Generator, Goyal et al., Nanda et al., lottery tickets, connectivity. The review retains exact primary passages, versions and raw hashes. Only small synthetic/bundled CPU contracts were executed: not an MNIST or Fashion-MNIST training result.