1import os, json, math, random
  2import numpy as np
  3import torch
  4from sklearn.datasets import load_digits
  5from sklearn.model_selection import train_test_split
  6
  7SEED = 2063
  8
  9def seed(s=SEED):
 10    random.seed(s); np.random.seed(s); torch.manual_seed(s)
 11    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
 12
 13
 14def toy_check():
 15    # V_line = y^2/2 has a 1-D zero manifold in R^2: lambda=(2-1)/2=0.5.
 16    # V_point = (x^2+y^2)/2 is an isolated zero: lambda=2/2=1.
 17    n = 1001
 18    x = np.linspace(-4, 4, n); dx = x[1] - x[0]
 19    X, Y = np.meshgrid(x, x, indexing='ij')
 20    Vline = 0.5 * Y**2
 21    Vpoint = 0.5 * (X**2 + Y**2)
 22    betas = np.array([5., 8., 12., 20., 35., 60., 100., 180., 300.])
 23    Zline = np.array([(np.exp(-b*Vline).sum()*dx*dx) for b in betas])
 24    Zpoint = np.array([(np.exp(-b*Vpoint).sum()*dx*dx) for b in betas])
 25    # Omit beta=5 to reduce finite-box effects.
 26    sl = np.polyfit(np.log(betas[1:]), np.log(Zline[1:]), 1)[0]
 27    sp = np.polyfit(np.log(betas[1:]), np.log(Zpoint[1:]), 1)[0]
 28    ratio_slope = np.polyfit(np.log(betas[1:]), np.log((Zline/Zpoint)[1:]), 1)[0]
 29    widths = []
 30    for b in betas:
 31        w = np.exp(-b*Vline)
 32        widths.append(math.sqrt(float((w * Y**2).sum() / w.sum())))
 33    width_slope = np.polyfit(np.log(betas[1:]), np.log(widths[1:]), 1)[0]
 34    return {
 35        'lambda_scaling': {
 36            'line_predicted': 0.5, 'line_observed': float(-sl),
 37            'point_predicted': 1.0, 'point_observed': float(-sp),
 38            'ratio_exponent_predicted': 0.5, 'ratio_exponent_observed': float(ratio_slope)},
 39        'width_scaling': {'predicted_exponent': -0.5, 'observed_exponent': float(width_slope)},
 40        'passed': bool(abs(-sl-.5)<.06 and abs(-sp-1.)<.08 and abs(ratio_slope-.5)<.08 and abs(width_slope+.5)<.06)
 41    }
 42
 43class MLP(torch.nn.Module):
 44    def __init__(self):
 45        super().__init__(); self.net=torch.nn.Sequential(torch.nn.Linear(64,96),torch.nn.ReLU(),torch.nn.Linear(96,10))
 46    def forward(self,x): return self.net(x)
 47
 48def loss_acc(model, X, y):
 49    with torch.no_grad():
 50        z=model(X); return float(torch.nn.functional.cross_entropy(z,y)), float((z.argmax(1)==y).float().mean())
 51
 52def perturb_loss(model, X, y, sigma=.02, reps=12):
 53    vals=[]; params=[p for p in model.parameters() if p.requires_grad]
 54    with torch.no_grad():
 55        original=[p.detach().clone() for p in params]
 56        for _ in range(reps):
 57            for p in params: p.add_(sigma*torch.randn_like(p))
 58            vals.append(loss_acc(model,X,y)[0])
 59            for p,o in zip(params,original): p.copy_(o)
 60    return float(np.mean(vals))
 61
 62def train_compare(device):
 63    data=load_digits(); X=data.data.astype('float32')/16.; y=data.target.astype('int64')
 64    Xtr,Xte,ytr,yte=train_test_split(X,y,test_size=.25,random_state=SEED,stratify=y)
 65    Xt=torch.tensor(Xtr,device=device); yt=torch.tensor(ytr,device=device)
 66    Xe=torch.tensor(Xte,device=device); ye=torch.tensor(yte,device=device)
 67    results={}
 68    for method in ['adamw','zero_langevin']:
 69        seed(); m=MLP().to(device)
 70        opt=torch.optim.AdamW(m.parameters(),lr=2e-3,weight_decay=1e-4)
 71        # Same ordinary training budget.
 72        for step in range(350):
 73            opt.zero_grad(); L=torch.nn.functional.cross_entropy(m(Xt),yt); L.backward(); opt.step()
 74        pre_train,pre_test=loss_acc(m,Xt,yt),loss_acc(m,Xe,ye)
 75        if method=='zero_langevin':
 76            eps=0.02; eta=2e-5; bmin,bmax=0.5,8.
 77            for step in range(500):
 78                beta=bmin*(bmax/bmin)**min(step/250.,1.)
 79                m.zero_grad(); raw=torch.nn.functional.cross_entropy(m(Xt),yt)
 80                # max(raw-eps,0) has the intended zero-loss stratum gradient.
 81                objective=torch.relu(raw-eps); objective.backward()
 82                with torch.no_grad():
 83                    for p in m.parameters():
 84                        p.add_(-eta*beta*p.grad + math.sqrt(2*eta)*torch.randn_like(p))
 85            phase='500 Langevin steps'
 86        else: phase='none'
 87        tr=loss_acc(m,Xt,yt); te=loss_acc(m,Xe,ye)
 88        results[method]={'pre_phase_train_loss':pre_train[0],'pre_phase_test_loss':pre_test[0],
 89                         'train_loss':tr[0],'train_acc':tr[1],'test_loss':te[0],'test_acc':te[1],
 90                         'perturbed_train_loss':perturb_loss(m,Xt,yt), 'phase':phase}
 91    return results
 92
 93def main():
 94    seed()
 95    try: device=torch.device('cuda' if torch.cuda.is_available() else 'cpu'); toy=toy_check(); exp=train_compare(device)
 96    except Exception as e:
 97        device=torch.device('cpu'); seed(); toy=toy_check(); exp=train_compare(device)
 98        exp['fallback_error']=repr(e)
 99    result={'device':str(device),'toy':toy,'experiment':exp}
100    os.makedirs('artifacts',exist_ok=True)
101    with open('artifacts/results.json','w') as f: json.dump(result,f,indent=2)
102    print(json.dumps(result,indent=2))
103
104if __name__=='__main__': main()