1 files changed, 197 insertions, 0 deletions
diff --git a/protocol/examples/audit_cnn.py b/protocol/examples/audit_cnn.py
new file mode 100644
index 0000000..890c1ec
--- /dev/null
+++ b/protocol/examples/audit_cnn.py
@@ -0,0 +1,197 @@
+"""
+Apply the protocol's diagnostic logic to the SmallCNN architecture (3 conv
+blocks + 1 FC + head, BatchNorm, no terminal LayerNorm). The existing
+checkpoints are in `results/cnn_baseline/{method}_s{seed}.pt`.
+
+This is a custom audit script (not via `protocol.diagnose(...)`) because
+the CNN has 4D conv hidden states and no `model.embed` / `model.out_ln`
+attributes that the duck-typed protocol API expects. The diagnostic
+*logic* is identical: per-block growth of flattened ‖h_l‖, BP grad floor
+at the deepest hidden layer, frozen-blocks comparison.
+
+Why this matters: CNN with BatchNorm is a third architecture family
+(neither pre-LN ResMLP nor pre-LN ViT). Both BP and DFA should be
+informative test cases:
+  - BP on CNN: should pass all diagnostics (sanity)
+  - DFA on CNN: open question — BatchNorm normalizes per-feature, so the
+    LN-driven gradient collapse mechanism may or may not apply
+
+Run:
+    CUDA_VISIBLE_DEVICES=2 python -m protocol.examples.audit_cnn
+"""
+import os
+import sys
+import json
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torchvision
+import torchvision.transforms as transforms
+from torch.utils.data import DataLoader
+
+REPO_ROOT = os.path.dirname(
+    os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+)
+sys.path.insert(0, REPO_ROOT)
+sys.path.insert(0, os.path.join(REPO_ROOT, "experiments"))
+
+# Import the SmallCNN from the experiments script
+import importlib.util
+_spec = importlib.util.spec_from_file_location(
+    "cnn_baseline_module",
+    os.path.join(REPO_ROOT, "experiments/cnn_baseline.py"),
+)
+_mod = importlib.util.module_from_spec(_spec)
+_spec.loader.exec_module(_mod)
+SmallCNN = _mod.SmallCNN
+
+
+CKPT_DIR = os.path.join(REPO_ROOT, "results/cnn_baseline")
+THRESHOLD_PER_BLOCK = 50.0
+THRESHOLD_GFLOOR = 1e-7
+
+
+def get_eval(n=1024, batch_size=128, device="cuda:0"):
+    tv = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),
+    ])
+    te = torchvision.datasets.CIFAR10("./data", train=False, download=True, transform=tv)
+    loader = DataLoader(te, batch_size=batch_size, shuffle=False, num_workers=0)
+    batches = []
+    for x, y in loader:
+        x, y = x.to(device), y.to(device)
+        batches.append((x, y))
+        if sum(b[0].size(0) for b in batches) >= n:
+            break
+    return batches
+
+
+def evaluate(model, device):
+    tv = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),
+    ])
+    te = torchvision.datasets.CIFAR10("./data", train=False, download=True, transform=tv)
+    loader = DataLoader(te, batch_size=256, shuffle=False, num_workers=0)
+    model.eval()
+    correct = total = 0
+    with torch.no_grad():
+        for x, y in loader:
+            x, y = x.to(device), y.to(device)
+            preds = model(x).argmax(-1)
+            correct += (preds == y).sum().item()
+            total += x.size(0)
+    return correct / total
+
+
+def per_layer_norms_and_grads(model, x, y):
+    """For the CNN, return per-layer flattened ‖h_l‖ medians and ‖g_l‖ medians."""
+    model.eval()
+    with torch.enable_grad():
+        h0 = model.blocks[0](x)
+        h1 = model.blocks[1](h0)
+        h2 = model.blocks[2](h1)
+        h3 = model.blocks[3](h2.flatten(1))
+        logits = model.out_head(h3)
+        hiddens = [h0, h1, h2, h3]
+        loss = F.cross_entropy(logits, y)
+        grads = torch.autograd.grad(loss, hiddens)
+
+    h_norms = []
+    g_norms = []
+    for h, g in zip(hiddens, grads):
+        h_flat = h.reshape(h.shape[0], -1)
+        g_flat = g.reshape(g.shape[0], -1)
+        h_norms.append(h_flat.norm(dim=-1).median().item())
+        g_norms.append(g_flat.norm(dim=-1).median().item())
+    return h_norms, g_norms
+
+
+def max_per_block_growth(h):
+    if len(h) < 2:
+        return 1.0
+    return max(h[i + 1] / max(h[i], 1e-30) for i in range(len(h) - 1))
+
+
+def load_cnn(method, seed, device):
+    path = os.path.join(CKPT_DIR, f"{method}_s{seed}.pt")
+    sd = torch.load(path, map_location=device, weights_only=False)
+    if isinstance(sd, dict) and "model_state" in sd:
+        sd = sd["model_state"]
+    elif isinstance(sd, dict) and "state_dict" in sd:
+        sd = sd["state_dict"]
+    model = SmallCNN().to(device)
+    model.load_state_dict(sd)
+    return model
+
+
+def main():
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    print(f"Device: {device}")
+    eval_batches = get_eval(n=1024, batch_size=128, device=device)
+    x, y = eval_batches[0]
+
+    methods = ["bp", "dfa", "state_bridge", "credit_bridge", "ep"]
+    print()
+    print("=" * 100)
+    print("CNN audit (SmallCNN: 3 conv + BN + 1 FC, NO terminal LN, CIFAR-10)")
+    print("=" * 100)
+    print(f"  {'method':<16}{'seed':>6}{'acc':>8}{'h_max/h_min':>14}{'max/block':>14}{'||g_L||':>14}  verdict")
+    print("  " + "-" * 100)
+
+    rows = []
+    for seed in [42, 123, 456]:
+        for method in methods:
+            try:
+                model = load_cnn(method, seed, device)
+            except Exception as e:
+                print(f"  {method:<16}{seed:>6}  SKIPPED ({e})")
+                continue
+            acc = evaluate(model, device)
+            h_norms, g_norms = per_layer_norms_and_grads(model, x, y)
+            max_growth = max_per_block_growth(h_norms)
+            h_ratio = max(h_norms) / max(min(h_norms), 1e-30)
+            g_L = g_norms[-1]
+            flags = []
+            if max_growth > THRESHOLD_PER_BLOCK:
+                flags.append("(a)")
+            if g_L < THRESHOLD_GFLOOR:
+                flags.append("(b)")
+            verdict = "trustworthy" if not flags else f"walk-back: {'+'.join(flags)}"
+            rows.append({
+                "method": method,
+                "seed": seed,
+                "acc": acc,
+                "h_norms": h_norms,
+                "g_norms": g_norms,
+                "max_per_block": max_growth,
+                "verdict": verdict,
+            })
+            print(f"  {method:<16}{seed:>6}{acc:>8.4f}{h_ratio:>14.2e}{max_growth:>14.2e}{g_L:>14.2e}  {verdict}")
+
+    print()
+    print("=" * 100)
+    print("Per-method 3-seed mean (h_norms across all 4 hidden layers, g across all):")
+    print("=" * 100)
+    for method in methods:
+        method_rows = [r for r in rows if r["method"] == method]
+        if not method_rows:
+            continue
+        accs = np.array([r["acc"] for r in method_rows])
+        h_arrs = np.array([r["h_norms"] for r in method_rows])
+        g_arrs = np.array([r["g_norms"] for r in method_rows])
+        max_g = np.array([r["max_per_block"] for r in method_rows])
+        print(f"  {method.upper()}: acc={accs.mean():.4f}±{accs.std():.4f}, "
+              f"h_means={h_arrs.mean(0)}, g_means={g_arrs.mean(0)}, "
+              f"max-per-block={max_g.mean():.2e}")
+
+    out_path = os.path.join(REPO_ROOT, "results/protocol_audit/audit_cnn_3seed.json")
+    with open(out_path, "w") as f:
+        json.dump(rows, f, indent=2)
+    print(f"\nSaved {out_path}")
+
+
+if __name__ == "__main__":
+    main()