1import sys, json
 2import numpy as np
 3import torch
 4import torch.nn as nn
 5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 6from bench import get_dataset, train_model, evaluate, sweep_baseline, make_report
 7from bench.models import cnn_small
 8
 9SEEDS = tuple(range(8)); EPOCHS = 8
10GRID = [{'lr': x, 'weight_decay': 0.0} for x in (0.0015, 0.003, 0.006)]
11
12class FeatureNormCNN(nn.Module):
13    # Identical cnn_small, differing only in the channel operation before its head.
14    def __init__(self, out_dim, mode='bn', rho=.8, eps=1e-4):
15        super().__init__(); base = cnn_small(out_dim)
16        self.features = base.net[:9]       # through third ReLU, before Flatten
17        self.tail = base.net[9:]           # exact shared Flatten/MLP/classifier
18        self.mode, self.rho, self.eps = mode, rho, eps
19        if mode == 'bn': self.norm = nn.BatchNorm1d(96)
20        else:
21            self.register_buffer('ema_R', torch.eye(96)); self.register_buffer('ema_mu', torch.zeros(96))
22            self.register_buffer('initialized', torch.tensor(False))
23
24    def _zca(self, z):
25        if self.training:
26            mu = z.mean(0); q = z - mu
27            R = q.T @ q / max(1, z.shape[0]) + self.eps*torch.eye(96, device=z.device)
28            with torch.no_grad():
29                if not bool(self.initialized): self.ema_R.copy_(R.detach()); self.ema_mu.copy_(mu.detach()); self.initialized.fill_(True)
30                else: self.ema_R.mul_(.95).add_(R.detach(), alpha=.05); self.ema_mu.mul_(.95).add_(mu.detach(), alpha=.05)
31        else: mu, R = self.ema_mu, self.ema_R + self.eps*torch.eye(96, device=z.device)
32        w,U = torch.linalg.eigh(R); A = (U*torch.rsqrt(torch.clamp(w,min=self.eps))) @ U.T
33        return (z-mu) @ A
34
35    def forward(self, x):
36        h = self.features(x); b,c,hh,ww = h.shape
37        z = h.permute(0,2,3,1).reshape(-1,c)
38        z = self.norm(z) if self.mode == 'bn' else self._zca(z)
39        h = z.reshape(b,hh,ww,c).permute(0,3,1,2)
40        return self.tail(h)
41
42def seed_all(s):
43    np.random.seed(s); torch.manual_seed(s)
44    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
45
46def train_fn(mode, cfg):
47    def run(seed):
48        seed_all(seed); d=get_dataset('vision', seed, n_train=400, n_test=200)
49        net=FeatureNormCNN(d['out_dim'], mode=mode, rho=cfg.get('rho',.8))
50        _,metric,_=train_model(net,d,epochs=EPOCHS,lr=cfg['lr'],weight_decay=cfg['weight_decay'],batch=128,log=lambda *_:None)
51        return metric
52    return run
53
54def signature(cfg):
55    seed_all(0); d=get_dataset('vision',0,n_train=400,n_test=200)
56    net=FeatureNormCNN(d['out_dim'],mode='zca',rho=cfg['rho']); net,_,_=train_model(net,d,epochs=EPOCHS,lr=cfg['lr'],weight_decay=cfg['weight_decay'],batch=128,log=lambda *_:None)
57    net.eval()
58    with torch.no_grad():
59        h=net.features(d['xte']); z=h.permute(0,2,3,1).reshape(-1,96); y=net._zca(z)
60        C=y.T@y/y.shape[0]; off=C-torch.diag(torch.diag(C))
61        zs=(z-z.mean(0))/(z.std(0)+1e-6); ys=(y-y.mean(0))/(y.std(0)+1e-6)
62        fid=torch.diag(zs.T@ys/zs.shape[0])
63    observed_off=float(torch.linalg.norm(off)); observed_f=float(fid.min())
64    return {'prediction':'decorrelation near zero with designated-channel fidelity >= rho_min', 'predicted_offdiag':0.0, 'predicted_min_fidelity_threshold':cfg['rho'], 'observed_offdiag':observed_off, 'observed_min_fidelity':observed_f, 'confirmed':bool(observed_off < .15 and observed_f >= cfg['rho']-.05)}
65
66def main():
67    base=sweep_baseline(lambda c:train_fn('bn',c),GRID,seeds=(0,1,2,3))
68    idea_grid=[dict(c,rho=.8) for c in GRID]
69    idea_runs=[]
70    for c in idea_grid: idea_runs.append((c,evaluate(train_fn('zca',c),seeds=SEEDS)))
71    best_cfg,best=max(idea_runs,key=lambda x:x[1]['mean']) if False else min(idea_runs,key=lambda x:x[1]['mean'])
72    rep=make_report('vision','cnn_small',base,best[1],extra=signature(best_cfg))
73    rep['idea_sweep']=[{'cfg':c,'result':r} for c,r in idea_runs]
74    rep['protocol_notes']='Vision chosen because anchored whitening is a channel normalization mechanism; same CNN and training budget used for BN and ZCA.'
75    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
76    print(json.dumps(rep,indent=2))
77if __name__=='__main__': main()