MNIST MLP vs CNN architecture comparison

# MNIST: MLP vs CNN head-to-head Shared scaffold below (data + model defs + one parameterized `train_and_eval`), then each architecture trains on **its own branch** for 3 epochs on the identical train/test split, followed by a side-by-side comparison.

import os, json, time, random
import numpy as np
import torch, torch.nn as nn, torch.nn.functional as F
from torch.utils.data import DataLoader
import torchvision
from torchvision import transforms

print("torch", torch.__version__, "| torchvision", torchvision.__version__)

SEED = 42
def seed_all(s=SEED):
    random.seed(s); np.random.seed(s); torch.manual_seed(s); torch.cuda.manual_seed_all(s)

seed_all()
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print("device:", DEVICE)

tf = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
…

torch 2.7.1+cpu | torchvision 0.22.1+cpu
device: cpu

from sklearn.metrics import confusion_matrix

BATCH_TRAIN, BATCH_TEST, EPOCHS, LR = 128, 512, 3, 1e-3

class MLP(nn.Module):
    """Baseline: flatten pixels -> 2 hidden dense layers."""
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Flatten(), nn.Linear(28*28, 256), nn.ReLU(),
            nn.Linear(256, 128), nn.ReLU(), nn.Linear(128, 10))
    def forward(self, x): return self.net(x)

class CNN(nn.Module):
    """Small conv net: 2 conv blocks + pooling -> dense head."""
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
…

scaffold ready: ['mlp', 'cnn'] | epochs = 3

import json, numpy as np, pandas as pd
import matplotlib.pyplot as plt

# Metrics collected from the two experiment branches (mlp-baseline / small-cnn)
R = {
 "mlp": json.loads(r'''{"arch": "mlp", "n_params": 235146, "final_test_acc": 0.969, "final_test_loss": 0.09835527067184448, "history": {"train_loss": [0.27202799055576327, 0.1023455897708734, 0.06919143068194389], "test_loss": [0.12825794093608855, 0.08895642035007477, 0.09835527067184448], "test_acc": [0.9613, 0.9715, 0.969], "epoch_time_s": [10.28, 9.61, 9.45]}, "confusion_matrix": [[970,0,0,2,0,2,1,3,1,1],[0,1128,0,1,0,1,2,1,2,0],[5,3,1003,5,1,1,1,13,0,0],[0,0,4,990,0,3,0,7,0,6],[0,0,6,0,943,0,2,4,0,27],[4,0,0,12,2,864,2,1,1,6],[7,3,1,1,8,14,917,4,3,0],[1,1,8,5,1,0,0,1009,0,3],[6,0,5,28,5,13,1,11,897,8],[3,2,0,8,4,2,0,21,0,969]], "per_class_errors": {"0":10,"1":7,"2":29,"3":20,"4":39,"5":28,"6":41,"7":19,"8":77,"9":40}}'''),
 "cnn": json.loads(r'''{"arch": "cnn", "n_params": 206922, "final_test_acc": 0.9864, "final_test_loss": 0.040551124787330625, "history": {"train_loss": [0.22926587956547737, 0.06047722478012244, 0.04254846897919973], "test_loss": [0.07385420913696289, 0.044217693829536435, 0.040551124787330625], "test_acc": [0.9783, 0.986, 0.9864], "epoch_time_s": [15.66, 14.84, 15.19]}, "confusion_matrix": [[975,0,0,0,0,0,3,1,1,0],[0,1124,2,3,0,0,2,1,3,0],[0,0,1021,1,0,0,0,7,3,0],[0,0,0,1002,0,4,0,1,1,2],[0,0,0,0,962,0,1,0,1,18],[3,0,0,8,0,873,2,1,1,4],[3,2,0,1,1,4,945,0,2,0],[0,1,8,2,0,0,0,1013,1,3],[4,0,3,4,2,1,1,2,952,5],[1,0,0,0,1,3,0,6,1,997]], "per_class_errors": {"0":5,"1":11,"2":11,"3":8,"4":20,"5":19,"6":13,"7":15,"8":22,"9":12}}'''),
}

summary = pd.DataFrame([{
    "model": k.upper(), "params": v["n_params"], "test_acc": v["final_test_acc"],
    "test_err_%": round(100*(1-v["final_test_acc"]), 2), "test_loss": round(v["final_test_loss"], 4),
    "total_errors": sum(v["per_class_errors"].values()),
    "sec/epoch": round(np.mean(v["history"]["epoch_time_s"]), 1),
} for k, v in R.items()]).set_index("model")
print(summary.to_string())
err_red = 100*(1 - (1-R["cnn"]["final_test_acc"])/(1-R["mlp"]["final_test_acc"]))
print(f"\nCNN cuts MNIST test error by {err_red:.1f}% relative to the MLP "
…

params  test_acc  test_err_%  test_loss  total_errors  sec/epoch
model                                                                  
MLP    235146    0.9690        3.10     0.0984           310        9.8
CNN    206922    0.9864        1.36     0.0406           136       15.2

CNN cuts MNIST test error by 56.1% relative to the MLP (310 -> 136 mistakes) with fewer parameters.

## Head-to-head findings (3 epochs, identical split, seed 42, Adam lr=1e-3) | | MLP (784-256-128-10) | Small CNN (16→32 conv + dense) | |---|---|---| | Params | 235,146 | **206,922** (fewer) | | Test accuracy | 0.9690 | **0.9864** | | Test loss | 0.0984 | **0.0406** | | Test mistakes / 10,000 | 310 | **136** | | Sec/epoch (CPU) | **9.8** | 15.2 | **Where they differ** 1. **Error rate:** the CNN removes 56% of the MLP's residual error (310 → 136 mistakes) while using ~12% fewer parameters — capacity is not the differentiator, inductive bias is. 2. **Generalisation trend:** the MLP's test loss *rose* in epoch 3 (0.0890 → 0.0984) while its train loss kept falling (0.102 → 0.069) — the start of overfitting. The CNN's train and test loss fall together (0.060 → 0.043 train, 0.0442 → 0.0406 test), so it was still improving when we stopped. 3. **Per-class:** the MLP's errors are concentrated on the topologically hard digits — 8 (77 errors), 6 (41), 9 (40), 4 (39). The CNN flattens this profile (worst class 8 with 22), i.e. it does not just improve on average, it removes the *specific* failure modes. 4. **Confusion pairs:** the biggest gains are 8→3 (28→4), 9→7 (21→6), 8→5 (13→1), 6→5 (14→4), 4→9 (27→18) — exactly the pairs distinguished by a small local stroke (a closed vs open loop, a crossbar) rather than by global pixel intensity. **Why** The MLP flattens the image, so every pixel is an independent input feature and spatial neighbourhood is destroyed; it must learn each strok…