%matplotlib inline
import json, time
from pathlib import Path
import numpy as np
import torch, torch.nn as nn, torch.nn.functional as F
import matplotlib.pyplot as plt
try:
    import torchvision
except ImportError:
    import subprocess, sys
    subprocess.run([sys.executable, "-m", "pip", "install", "-q", "torchvision"], check=True)
    import torchvision
from torchvision import datasets

TRAIN_SIZE = 1000
STEPS      = 40000
EVAL_EVERY = 200
LR         = 1e-3
SEED       = 0
WIDTH      = 256

FIG_DIR = Path("figures"); FIG_DIR.mkdir(exist_ok=True)
RES_DIR = Path("results"); RES_DIR.mkdir(exist_ok=True)
device = "cuda" if torch.cuda.is_available() else "cpu"
print("device:", device)

device: cuda

_tr = datasets.MNIST("data", train=True,  download=True)
_te = datasets.MNIST("data", train=False, download=True)
g = torch.Generator().manual_seed(SEED)
idx = torch.randperm(len(_tr.data), generator=g)[:TRAIN_SIZE]
ytr = _tr.targets[idx].to(device); yte = _te.targets.to(device)
Ytr = F.one_hot(ytr, 10).float()
Xtr_flat = (_tr.data[idx].float()/255.).reshape(-1,784).to(device)
Xte_flat = (_te.data.float()/255.).reshape(-1,784).to(device)
Xtr_img  = (_tr.data[idx].float()/255.).unsqueeze(1).to(device)
Xte_img  = (_te.data.float()/255.).unsqueeze(1).to(device)
print("train:", TRAIN_SIZE, "test:", Xte_flat.shape[0])

train: 1000 test: 10000

class MLP(nn.Module):
    def __init__(self, width=256):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(784,width), nn.ReLU(),
                                 nn.Linear(width,width), nn.ReLU(),
                                 nn.Linear(width,10))
    def forward(self,x): return self.net(x)

class CNN(nn.Module):  # BatchNorm-free (BN would cancel the init-scale trick)
    def __init__(self):
        super().__init__()
        self.f = nn.Sequential(nn.Conv2d(1,32,3,padding=1), nn.ReLU(), nn.MaxPool2d(2),
                               nn.Conv2d(32,64,3,padding=1), nn.ReLU(), nn.MaxPool2d(2))
        self.c = nn.Sequential(nn.Flatten(), nn.Linear(64*7*7,128), nn.ReLU(), nn.Linear(128,10))
    def forward(self,x): return self.c(self.f(x))

def build(arch, alpha):
    torch.manual_seed(SEED)
    m = (MLP(WIDTH) if arch=="MLP" else CNN()).to(device)
    with torch.no_grad():
        for p in m.parameters(): p.mul_(alpha)
    return m

def views(arch):
    return (Xtr_flat, Xte_flat) if arch=="MLP" else (Xtr_img, Xte_img)

@torch.no_grad()
def acc(model, X, y, chunk=5000):
    model.eval(); c=n=0
    for i in range(0, X.shape[0], chunk):
        c += (model(X[i:i+chunk]).argmax(-1)==y[i:i+chunk]).sum().item(); n += X[i:i+chunk].shape[0]
    return c/n

def train_run(arch, loss, alpha, wd):
    model = build(arch, alpha)
    opt = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=wd)
    Xtr, Xte = views(arch)
    hist = {"step": [], "train_acc": [], "test_acc": [], "wnorm": []}
    t0 = time.time()
    for step in range(STEPS+1):
        model.train()
        out = model(Xtr)
        L = F.mse_loss(out, Ytr) if loss=="mse" else F.cross_entropy(out, ytr)
        opt.zero_grad(); L.backward(); opt.step()
        if step % EVAL_EVERY == 0:
            tr = acc(model, Xtr, ytr); te = acc(model, Xte, yte)
            wn = sum(p.pow(2).sum() for p in model.parameters()).sqrt().item()
            for k,v in zip(hist, [step, tr, te, wn]): hist[k].append(v)
            if step % (EVAL_EVERY*40) == 0:
                print(f"[{arch}/{loss}/a{alpha}/wd{wd}] step {step:5d} tr {tr:.3f} te {te:.3f} wn {wn:.0f} ({time.time()-t0:.0f}s)", flush=True)
    return hist

def mem_step(h, thr=0.99):
    for s,v in zip(h["step"], h["train_acc"]):
        if v>=thr: return s
    return None
def test_at_mem(h):
    s = mem_step(h)
    return h["test_acc"][h["step"].index(s)] if s is not None else h["test_acc"][-1]

configs, seen = [], set()
def add(arch, loss, alpha, wd):
    k=(arch,loss,alpha,wd)
    if k not in seen: seen.add(k); configs.append(k)
for arch in ["MLP","CNN"]:
    for loss in ["mse","ce"]:
        for alpha in [1,8]:
            add(arch,loss,alpha,0.1)
add("MLP","mse",8,0.0)
for loss in ["mse","ce"]:
    for alpha in [1,2,4,8,16]:
        add("MLP",loss,alpha,0.1)

def key(arch,loss,alpha,wd): return f"{arch}_{loss}_a{alpha}_wd{wd}"

results = {}
for (arch,loss,alpha,wd) in configs:
    print(f"=== {arch} {loss} alpha={alpha} wd={wd} ===", flush=True)
    results[key(arch,loss,alpha,wd)] = train_run(arch, loss, alpha, wd)

with open(RES_DIR/"ablation_metrics.json","w") as f:
    json.dump({"config":{"TRAIN_SIZE":TRAIN_SIZE,"STEPS":STEPS,"LR":LR,"SEED":SEED},
               "results":results}, f)
print("saved results/ablation_metrics.json", flush=True)

=== MLP mse alpha=1 wd=0.1 ===

[MLP/mse/a1/wd0.1] step     0 tr 0.363 te 0.350 wn 13 (0s)

[MLP/mse/a1/wd0.1] step  8000 tr 1.000 te 0.923 wn 15 (19s)

[MLP/mse/a1/wd0.1] step 16000 tr 1.000 te 0.924 wn 18 (37s)

[MLP/mse/a1/wd0.1] step 24000 tr 1.000 te 0.926 wn 21 (55s)

[MLP/mse/a1/wd0.1] step 32000 tr 1.000 te 0.926 wn 20 (74s)

[MLP/mse/a1/wd0.1] step 40000 tr 1.000 te 0.929 wn 22 (93s)

=== MLP mse alpha=8 wd=0.1 ===

[MLP/mse/a8/wd0.1] step     0 tr 0.072 te 0.063 wn 106 (0s)

[MLP/mse/a8/wd0.1] step  8000 tr 1.000 te 0.755 wn 49 (19s)

def xs(h):
    s=np.array(h["step"]); s[0]=1; return s

# 1) knockout grid: 2x2 (arch x loss) at alpha=8, wd=0.1; overlay wd=0 trap in MLP/mse
fig, axes = plt.subplots(2,2, figsize=(11,8), sharex=True, sharey=True)
for i,arch in enumerate(["MLP","CNN"]):
    for j,loss in enumerate(["mse","ce"]):
        ax=axes[i][j]; h=results[key(arch,loss,8,0.1)]
        ax.plot(xs(h), h["train_acc"], color="tab:blue", label="train")
        ax.plot(xs(h), h["test_acc"],  color="tab:red",  label="test")
        if arch=="MLP" and loss=="mse":
            h0=results[key("MLP","mse",8,0.0)]
            ax.plot(xs(h0), h0["test_acc"], color="tab:gray", ls="--", label="test (wd=0)")
        ax.set_xscale("log"); ax.set_ylim(-0.02,1.02); ax.grid(alpha=0.3)
        ax.set_title(arch+" + "+loss.upper()+"  (alpha=8, wd=0.1)")
        if i==1: ax.set_xlabel("optimizer step")
        if j==0: ax.set_ylabel("accuracy")
axes[0][0].legend(fontsize=8)
fig.suptitle("Grokking with MSE + large init + weight decay — both MLP and CNN; CE never groks")
fig.tight_layout(); fig.savefig(FIG_DIR/"ablation_curves.png", dpi=130); plt.show()

# 2) alpha sweep (MLP): test accuracy AT MEMORIZATION vs alpha, MSE vs CE
ALPHAS=[1,2,4,8,16]
fig, ax = plt.subplots(figsize=(6.5,4.3))
for loss,c in [("mse","tab:red"),("ce","tab:blue")]:
    tam=[test_at_mem(results[key("MLP",loss,a,0.1)]) for a in ALPHAS]
    fin=[results[key("MLP",loss,a,0.1)]["test_acc"][-1] for a in ALPHAS]
    ax.plot(ALPHAS, tam, "o-", color=c, label=loss.upper()+" (at memorization)")
    ax.plot(ALPHAS, fin, "o--", color=c, alpha=0.5, label=loss.upper()+" (final)")
ax.set_xscale("log", base=2); ax.set_xticks(ALPHAS); ax.set_xticklabels(ALPHAS)
ax.set_xlabel("init scale alpha"); ax.set_ylabel("test accuracy"); ax.set_ylim(-0.02,1.02); ax.grid(alpha=0.3)
ax.legend(fontsize=8); ax.set_title("MLP: MSE digs a memorization trap as init grows; CE does not")
fig.tight_layout(); fig.savefig(FIG_DIR/"ablation_alpha.png", dpi=130); plt.show()

# 3) mechanism: weight norm vs test accuracy (MLP/MSE/alpha=8/wd=0.1)
h=results[key("MLP","mse",8,0.1)]
fig, ax1 = plt.subplots(figsize=(6.8,4.3)); ax2=ax1.twinx()
ax1.plot(xs(h), h["wnorm"], color="tab:purple", label="weight norm")
ax2.plot(xs(h), h["test_acc"], color="tab:red", ls="--", label="test acc")
ax1.set_xscale("log"); ax1.set_xlabel("optimizer step")
ax1.set_ylabel("weight L2 norm (solid)"); ax2.set_ylabel("test accuracy (dashed)")
ax1.set_title("Mechanism: test rises as the norm migrates down (MLP/MSE/a=8/wd=0.1)")
fig.tight_layout(); fig.savefig(FIG_DIR/"ablation_norm.png", dpi=130); plt.show()

print("%-22s %8s %8s %8s" % ("config","mem@","test@mem","final"))
for (arch,loss,alpha,wd) in configs:
    h=results[key(arch,loss,alpha,wd)]
    print("%-22s %8s %8.3f %8.3f" % (key(arch,loss,alpha,wd), str(mem_step(h)), test_at_mem(h), h["test_acc"][-1]))
print("saved figures:", sorted(p.name for p in FIG_DIR.glob("ablation_*.png")))

config                     mem@ test@mem    final
MLP_mse_a1_wd0.1            200    0.917    0.929
MLP_mse_a8_wd0.1           1200    0.208    0.928
MLP_ce_a1_wd0.1             200    0.885    0.862
MLP_ce_a8_wd0.1             200    0.814    0.866
CNN_mse_a1_wd0.1            200    0.969    0.965
CNN_mse_a8_wd0.1           4800    0.327    0.964
CNN_ce_a1_wd0.1             200    0.937    0.924
CNN_ce_a8_wd0.1             200    0.855    0.922
MLP_mse_a8_wd0.0           1200    0.178    0.243
MLP_mse_a2_wd0.1            200    0.912    0.929
MLP_mse_a4_wd0.1            200    0.721    0.926
MLP_mse_a16_wd0.1          4000    0.127    0.924
MLP_ce_a2_wd0.1             200    0.897    0.856
MLP_ce_a4_wd0.1             200    0.900    0.889
MLP_ce_a16_wd0.1            200    0.740    0.864
saved figures: ['ablation_alpha.png', 'ablation_curves.png', 'ablation_norm.png']

Inducing grokking on natural (image) data — a loss x architecture x init ablation¶