1import json, math, random
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED = 471
  7
  8def seed_all(s):
  9    random.seed(s); np.random.seed(s); torch.manual_seed(s)
 10    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
 11
 12def relu_cov(K, variance_scale=2.0):
 13    # E[ReLU(u) ReLU(v)] for a zero-mean Gaussian pair, followed by He scale.
 14    d = torch.sqrt(torch.clamp(torch.diag(K), min=1e-12))
 15    corr = K / (d[:, None] * d[None, :])
 16    corr = torch.clamp(corr, -1.0, 1.0)
 17    theta = torch.acos(corr)
 18    e = (d[:, None] * d[None, :]) * (torch.sin(theta) + (math.pi-theta)*corr) / (2*math.pi)
 19    return variance_scale * e
 20
 21def nngp_targets(X, depth):
 22    K = X @ X.T / X.shape[1]
 23    Ks = []
 24    for _ in range(depth):
 25        K = relu_cov(K)
 26        Ks.append(K)
 27    return Ks
 28
 29class MLP(nn.Module):
 30    def __init__(self, d, width, depth):
 31        super().__init__()
 32        self.layers = nn.ModuleList()
 33        self.layers.append(nn.Linear(d, width, bias=False))
 34        for _ in range(depth-1): self.layers.append(nn.Linear(width, width, bias=False))
 35        self.out = nn.Linear(width, 1, bias=False)
 36        for m in self.layers:
 37            nn.init.normal_(m.weight, std=math.sqrt(2/m.in_features))
 38        nn.init.normal_(self.out.weight, std=1/math.sqrt(width))
 39    def forward(self, x):
 40        hs=[]
 41        h=x
 42        for layer in self.layers:
 43            h=torch.relu(layer(h)); hs.append(h)
 44        return self.out(h), hs
 45
 46def covariance(h):
 47    return h @ h.T / h.shape[1]
 48
 49def train(reg, device, seed):
 50    seed_all(seed)
 51    n,d=192,10
 52    X=torch.randn(n,d,device=device)
 53    # A nonlinear but learnable fixed task, with a held-out validation set.
 54    y=(torch.sin(X[:,0]*1.4)+0.35*X[:,1]*X[:,2]).unsqueeze(1)
 55    Xt=torch.randn(96,d,device=device)
 56    yt=(torch.sin(Xt[:,0]*1.4)+0.35*Xt[:,1]*Xt[:,2]).unsqueeze(1)
 57    width,depth=64,3; batch=32
 58    model=MLP(d,width,depth).to(device)
 59    opt=torch.optim.Adam(model.parameters(),lr=3e-3)
 60    # The NNGP target is computed from the same input minibatch, and detached.
 61    history=[]; cov_history=[]
 62    for step in range(240):
 63        ix=torch.arange((step*batch)%n, (step*batch)%n+batch, device=device)%n
 64        xb,yb=X[ix],y[ix]
 65        pred,hs=model(xb)
 66        task=((pred-yb)**2).mean()
 67        targets=nngp_targets(xb,depth)
 68        covloss=sum(((covariance(h)-k.detach())**2).mean() for h,k in zip(hs,targets))/depth
 69        loss=task + (0.03*covloss if reg else 0.0)
 70        opt.zero_grad(); loss.backward(); opt.step()
 71        if step in (0,39,119,239):
 72            with torch.no_grad():
 73                vp,vhs=model(Xt); val=((vp-yt)**2).mean().item()
 74                # Evaluate deviation on a fresh calibration batch, as proposed.
 75                _,eh=model(X[:batch]); tk=nngp_targets(X[:batch],depth)
 76                dev=float(sum(((covariance(h)-k)**2).mean() for h,k in zip(eh,tk))/depth)
 77            history.append(val); cov_history.append(dev)
 78    return history[-1], cov_history[-1], history, cov_history
 79
 80def scaling_check():
 81    seed_all(SEED)
 82    d,B=12,24
 83    X=torch.randn(B,d)
 84    K=nngp_targets(X,2)[-1]
 85    widths=[32,64,128,256,512]
 86    reps=40
 87    means=[]
 88    for w in widths:
 89        errs=[]
 90        for r in range(reps):
 91            # independent He network, no output layer; compare postactivation Gram.
 92            W=torch.randn(d,w)*math.sqrt(2/d)
 93            h=torch.relu(X@W)
 94            W2=torch.randn(w,w)*math.sqrt(2/w)
 95            h=torch.relu(h@W2)
 96            errs.append(torch.linalg.norm(covariance(h)-K).item())
 97        means.append(float(np.mean(errs)))
 98    slope=float(np.polyfit(np.log(widths),np.log(means),1)[0])
 99    # Also verify the scalar ReLU formula at correlation values by Monte Carlo.
100    z=torch.randn(800000,2)
101    rho=0.6; z[:,1]=rho*z[:,0]+math.sqrt(1-rho*rho)*z[:,1]
102    mc=(torch.relu(z[:,0])*torch.relu(z[:,1])).mean().item()*2
103    kk=torch.tensor([[1.,rho],[rho,1.]])
104    analytic=relu_cov(kk)[0,1].item()
105    return widths,means,slope,mc,analytic
106
107def main():
108    try: device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
109    except Exception: device=torch.device('cpu')
110    try:
111        scaling=scaling_check()
112        runs=[]
113        for s in [471,472,473]:
114            b=train(False,device,s); r=train(True,device,s)
115            runs.append({'seed':s,'baseline_val':b[0],'idea_val':r[0], 'baseline_cov':b[1], 'idea_cov':r[1], 'baseline_curve':b[2], 'idea_curve':r[2], 'baseline_cov_curve':b[3], 'idea_cov_curve':r[3]})
116        out={'device':str(device),'scaling':{'widths':scaling[0],'errors':scaling[1],'loglog_slope':scaling[2],'mc_relu_cov':scaling[3],'analytic_relu_cov':scaling[4]},'runs':runs}
117    except Exception as e:
118        if device.type=='cuda':
119            device=torch.device('cpu'); scaling=scaling_check(); runs=[]
120            for s in [471,472,473]:
121                b=train(False,device,s); r=train(True,device,s)
122                runs.append({'seed':s,'baseline_val':b[0],'idea_val':r[0], 'baseline_cov':b[1], 'idea_cov':r[1]})
123            out={'device':'cpu_fallback','scaling':{'widths':scaling[0],'errors':scaling[1],'loglog_slope':scaling[2],'mc_relu_cov':scaling[3],'analytic_relu_cov':scaling[4]},'runs':runs,'cuda_error':repr(e)}
124        else: raise
125    with open('results.json','w') as f: json.dump(out,f,indent=2)
126    print(json.dumps(out,indent=2))
127if __name__=='__main__': main()