Add BP supplement for A2 CIFAR: 10 seeds with acc, Gamma, rho, naive_StateErr

BP 10-seed results: acc=0.614±0.003, Gamma=1.0, rho=0.998
Appended to A2_cifar_state_vs_credit.csv and A2_naive_state_err.csv

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

1 files changed, 107 insertions, 0 deletions

diff --git a/experiments/a2_bp_supplement.py b/experiments/a2_bp_supplement.py
new file mode 100644
index 0000000..c594f9e
--- /dev/null
+++ b/experiments/a2_bp_supplement.py
@@ -0,0 +1,107 @@
+"""Supplement A2 with BP baseline: 10 seeds, CIFAR-10, L=4, d=256."""
+import os, sys, json, csv, argparse, numpy as np, torch, torch.nn as nn, torch.nn.functional as F
+import torch.optim as optim
+from torch.utils.data import DataLoader
+import torchvision, torchvision.transforms as transforms
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from models.residual_mlp import ResidualMLP
+from metrics.credit_metrics import cosine_similarity_batch, perturbation_correlation
+
+def get_cifar10(bs=128):
+    tt=transforms.Compose([transforms.RandomCrop(32,4),transforms.RandomHorizontalFlip(),transforms.ToTensor(),transforms.Normalize((0.4914,0.4822,0.4465),(0.2470,0.2435,0.2616))])
+    tv=transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.4914,0.4822,0.4465),(0.2470,0.2435,0.2616))])
+    return (DataLoader(torchvision.datasets.CIFAR10('./data',True,download=True,transform=tt),bs,True,num_workers=4,pin_memory=True),
+            DataLoader(torchvision.datasets.CIFAR10('./data',False,download=True,transform=tv),bs,False,num_workers=4,pin_memory=True))
+
+def evaluate(m,tl,dev):
+    m.eval();c,t=0,0
+    with torch.no_grad():
+        for x,y in tl:x=x.view(x.size(0),-1).to(dev);y=y.to(dev);c+=(m(x).argmax(1)==y).sum().item();t+=x.size(0)
+    return c/t
+
+def compute_naive_state_err(model, test_loader, device, eval_layer=None):
+    model.eval(); L=model.num_blocks
+    if eval_layer is None: eval_layer=L//2
+    te,n=0.0,0
+    with torch.no_grad():
+        for x,y in test_loader:
+            x=x.view(x.size(0),-1).to(device);_,hi=model(x,return_hidden=True)
+            hL=hi[-1];hl=hi[eval_layer];norm=hL.norm(-1,keepdim=True).clamp(min=1.0)
+            te+=((hl-hL)/norm).pow(2).sum(-1).mean().item()*x.size(0);n+=x.size(0)
+    return te/n
+
+def compute_bp_diagnostics(model, test_loader, device):
+    model.eval(); L=model.num_blocks
+    for x,y in test_loader:
+        x=x.view(x.size(0),-1).to(device);y=y.to(device);break
+    batch=x.size(0)
+    lo,hbp=model(x,return_hidden=True)
+    for l in range(L+1):hbp[l].retain_grad()
+    F.cross_entropy(lo,y).backward()
+    bp={l:hbp[l].grad.detach().clone() for l in range(L+1)}
+    # BP credit = BP grad, so Gamma=1.0 by definition, rho≈1.0
+    # But let's compute rho properly
+    with torch.no_grad():_,hi=model(x,return_hidden=True)
+    gammas,rhos=[],[]
+    for l in range(L):
+        gammas.append(1.0)  # BP cosine with itself
+        h_l=hi[l].detach();a_l=bp[l]
+        def make_fwd(sl):
+            def f(h):
+                with torch.no_grad():
+                    c=h
+                    for i in range(sl,L):c=c+model.blocks[i](c)
+                    return F.cross_entropy(model.out_head(model.out_ln(c)),y,reduction='none')
+            return f
+        rhos.append(perturbation_correlation(h_l,a_l,make_fwd(l),epsilon=1e-3,M=16))
+    return float(np.mean(gammas)),float(np.mean(rhos))
+
+def main():
+    p=argparse.ArgumentParser()
+    p.add_argument('--gpu',type=int,default=3)
+    p.add_argument('--output_dir',type=str,default='results/confirmatory')
+    a=p.parse_args()
+    device=torch.device(f'cuda:{a.gpu}')
+    os.makedirs(a.output_dir,exist_ok=True)
+    seeds=[42,123,456,789,1024,2048,3000,4000,5000,6000]
+    L,d=4,256;trl,tel=get_cifar10()
+    rows_main=[];rows_naive=[]
+    for seed in seeds:
+        torch.manual_seed(seed);np.random.seed(seed);torch.cuda.manual_seed_all(seed)
+        model=ResidualMLP(3072,d,10,L).to(device)
+        opt=optim.AdamW(model.parameters(),lr=1e-3,weight_decay=0.01)
+        sch=optim.lr_scheduler.CosineAnnealingLR(opt,T_max=100)
+        print(f"  BP seed={seed}: training...",flush=True)
+        for ep in range(1,101):
+            model.train()
+            for x,y in trl:
+                x=x.view(x.size(0),-1).to(device);y=y.to(device)
+                loss=F.cross_entropy(model(x),y);opt.zero_grad();loss.backward();opt.step()
+            sch.step()
+            if ep%20==0:print(f"    Ep {ep}: acc={evaluate(model,tel,device):.4f}",flush=True)
+        acc=evaluate(model,tel,device)
+        gamma,rho=compute_bp_diagnostics(model,tel,device)
+        nse=compute_naive_state_err(model,tel,device)
+        rows_main.append({'method':'bp','seed':seed,'StateErr':float('nan'),'Gamma':gamma,'rho':rho,'acc':acc})
+        rows_naive.append({'method':'bp','seed':seed,'naive_StateErr':nse})
+        print(f"  BP seed={seed}: acc={acc:.4f}, Gamma={gamma:.4f}, rho={rho:.4f}, naive_StateErr={nse:.6f}",flush=True)
+    # Append to A2 CSV
+    a2_path=os.path.join(a.output_dir,'A2_cifar_state_vs_credit.csv')
+    with open(a2_path,'a',newline='') as f:
+        w=csv.DictWriter(f,fieldnames=['method','seed','StateErr','Gamma','rho','acc']);w.writerows(rows_main)
+    print(f"Appended {len(rows_main)} BP rows to {a2_path}",flush=True)
+    # Append to naive CSV
+    naive_path=os.path.join(a.output_dir,'A2_naive_state_err.csv')
+    with open(naive_path,'a',newline='') as f:
+        w=csv.DictWriter(f,fieldnames=['method','seed','naive_StateErr']);w.writerows(rows_naive)
+    print(f"Appended {len(rows_naive)} BP rows to {naive_path}",flush=True)
+    # Also save standalone
+    bp_path=os.path.join(a.output_dir,'A2_bp_supplement.csv')
+    with open(bp_path,'w',newline='') as f:
+        w=csv.DictWriter(f,fieldnames=['method','seed','StateErr','Gamma','rho','acc','naive_StateErr'])
+        w.writeheader()
+        for m,n in zip(rows_main,rows_naive):
+            w.writerow({**m,'naive_StateErr':n['naive_StateErr']})
+    print(f"Saved standalone to {bp_path}",flush=True)
+
+if __name__=='__main__':main()