1import os, json, math, random, time
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5
  6SEED = 558
  7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  8torch.set_num_threads(min(12, os.cpu_count() or 1))
  9try:
 10    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 11    if device.type == 'cuda':
 12        torch.cuda.reset_peak_memory_stats()
 13except Exception:
 14    device = torch.device('cpu')
 15
 16
 17def L_min(n, d, m, eps, c0=0.25, C0=2.0, delta=1e-6):
 18    eps = max(float(eps), delta)
 19    q = (m + 1) * math.log(max(C0 * (m + 1) * n * d / eps, 1.000001))
 20    return c0 * eps * math.sqrt(n / q)
 21
 22class MLP(nn.Module):
 23    def __init__(self, d, width):
 24        super().__init__()
 25        self.d, self.width = d, width
 26        self.net = nn.Sequential(nn.Linear(d, width), nn.ReLU(), nn.Linear(width, 1))
 27    def forward(self, x): return self.net(x).squeeze(-1)
 28    def widen(self, new_width):
 29        old = self.net
 30        new = MLP(self.d, new_width).to(next(self.parameters()).device)
 31        with torch.no_grad():
 32            new.net[0].weight[:self.width].copy_(old[0].weight)
 33            new.net[0].bias[:self.width].copy_(old[0].bias)
 34            new.net[2].weight[:, :self.width].copy_(old[2].weight)
 35            new.net[2].bias.copy_(old[2].bias)
 36            # New units have small output coefficients, as prescribed.
 37            new.net[2].weight[:, self.width:].mul_(0.02)
 38        return new
 39
 40def jacobian_lip(model, x):
 41    # Exact scalar-output Jacobian row norm on a small probe set.
 42    vals = []
 43    for z in x:
 44        zz = z.detach().clone().requires_grad_(True)
 45        y = model(zz[None])[0]
 46        g = torch.autograd.grad(y, zz, retain_graph=False)[0]
 47        vals.append(g.norm().item())
 48    return max(vals) if vals else 0.0
 49
 50def run(controller, seed=SEED):
 51    torch.manual_seed(seed); np.random.seed(seed)
 52    n, d, noise = 320, 5, 0.35
 53    x = torch.randn(n, d, device=device)
 54    clean = torch.sin(x[:, 0]) + 0.35 * x[:, 1] - 0.2 * x[:, 2]**2
 55    y = clean + noise * torch.randn(n, device=device)
 56    perm = torch.randperm(n, device=device)
 57    tr, va = perm[:240], perm[240:]
 58    model = MLP(d, 8).to(device)
 59    opt = torch.optim.Adam(model.parameters(), lr=0.015)
 60    history=[]; widen_count=0; t0=time.time()
 61    # Same number of update intervals for both conditions; widening adds capacity, not extra steps.
 62    for step in range(1, 801):
 63        model.train(); opt.zero_grad()
 64        loss = ((model(x[tr]) - y[tr])**2).mean()
 65        loss.backward(); opt.step()
 66        if step % 80 == 0:
 67            model.eval()
 68            with torch.no_grad():
 69                train_mse = ((model(x[tr])-y[tr])**2).mean().item()
 70                val_noisy = ((model(x[va])-y[va])**2).mean().item()
 71                val_clean = ((model(x[va])-clean[va])**2).mean().item()
 72            # Probe sensitivity on held-out points. The theorem's constants are empirical,
 73            # deliberately explicit and fixed across both runs.
 74            lip = jacobian_lip(model, x[va[:32]])
 75            eps = math.sqrt(max(train_mse, 1e-12))
 76            bound = L_min(len(tr), d, model.width, eps)
 77            if controller and train_mse < noise**2 and lip < 0.9 * bound and model.width < 64:
 78                oldw=model.width; model=model.widen(min(64, max(oldw+1, int(math.ceil(oldw*1.5)))))
 79                opt=torch.optim.Adam(model.parameters(), lr=0.015)
 80                widen_count += 1
 81            history.append({'step':step,'width':model.width,'train_mse':train_mse,'val_noisy_mse':val_noisy,'val_clean_mse':val_clean,'lip':lip,'L_min':bound})
 82    # Noise robustness is measured against the clean target on fresh Gaussian probes.
 83    with torch.no_grad():
 84        xte=torch.randn(160,d,device=device)
 85        clean_te=torch.sin(xte[:,0])+0.35*xte[:,1]-0.2*xte[:,2]**2
 86        pred=model(xte)
 87        clean_mse=((pred-clean_te)**2).mean().item()
 88    return {'history':history,'final_width':model.width,'widen_count':widen_count,'test_clean_mse':clean_mse,'seconds':time.time()-t0}
 89
 90def math_check():
 91    # Claimed scaling: increasing m lowers the necessary lower bound, while increasing
 92    # epsilon raises it. Ratios are checked against monotonicity, not a fitted theorem.
 93    widths=[2,4,8,16,32,64]
 94    by_width=[L_min(240,5,m,0.35) for m in widths]
 95    errs=[0.08,0.16,0.32,0.64]
 96    by_error=[L_min(240,5,16,e) for e in errs]
 97    return {'widths':widths,'L_by_width':by_width,'errors':errs,'L_by_error':by_error,
 98            'width_monotone':all(a>b for a,b in zip(by_width,by_width[1:])),
 99            'error_monotone':all(a<b for a,b in zip(by_error,by_error[1:]))}
100
101def main():
102    mc=math_check()
103    base=run(False, SEED)
104    idea=run(True, SEED)
105    out={'device':str(device),'math_check':mc,'baseline':base,'idea':idea,
106         'constants':{'c0':0.25,'C0':2.0,'eta':0.9},
107         'interpretation':'Controller widens only when empirical error is below known noise variance and estimated Jacobian Lipschitz constant is below the bound.'}
108    with open('results.json','w') as f: json.dump(out,f,indent=2)
109    print(json.dumps({'device':str(device),'math_check':mc,
110        'baseline':{k:base[k] for k in ['final_width','widen_count','test_clean_mse','seconds']},
111        'idea':{k:idea[k] for k in ['final_width','widen_count','test_clean_mse','seconds']}}, indent=2))
112if __name__=='__main__': main()