1import json, math, random, time
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6
  7SEED = 3142
  8
  9def seed_all(seed=SEED):
 10    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 11    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 12
 13def device_or_cpu():
 14    if torch.cuda.is_available():
 15        try:
 16            x = torch.zeros(1, device='cuda'); del x
 17            return torch.device('cuda')
 18        except Exception as e:
 19            print('CUDA fallback:', repr(e))
 20    return torch.device('cpu')
 21
 22class PlainMLP(nn.Module):
 23    def __init__(self, d, width=64, depth=3):
 24        super().__init__()
 25        layers=[nn.Linear(d,width), nn.Tanh()]
 26        for _ in range(depth-1): layers += [nn.Linear(width,width), nn.Tanh()]
 27        layers += [nn.Linear(width,1)]
 28        self.net=nn.Sequential(*layers)
 29    def forward(self,x): return self.net(x).squeeze(-1)
 30
 31class AccumulatorResNet(nn.Module):
 32    """Persistent context h and additive scalar accumulator s.
 33    alpha=0 transmits h exactly; each branch adds one correction to s.
 34    """
 35    def __init__(self,d,q=32,r=1,K=4,width=48):
 36        super().__init__(); self.K=K; self.q=q; self.r=r
 37        self.embed=nn.Sequential(nn.Linear(d,q),nn.Tanh())
 38        self.s0=nn.Parameter(torch.zeros(r))
 39        self.branches=nn.ModuleList()
 40        for _ in range(K):
 41            self.branches.append(nn.Sequential(nn.Linear(q+r,width),nn.Tanh(),
 42                                               nn.Linear(width,q+r),nn.Tanh()))
 43        self.head=nn.Sequential(nn.Linear(q+r,width),nn.Tanh(),nn.Linear(width,1))
 44    def forward(self,x):
 45        h=self.embed(x); s=self.s0.expand(x.shape[0],-1)
 46        for branch in self.branches:
 47            dh,ds=branch(torch.cat([h,s],1)).split([self.q,self.r],1)
 48            # paper-motivated gates: context is carried, accumulator adds correction
 49            h = h + 0.0*dh
 50            s = s + ds
 51        return self.head(torch.cat([h,s],1)).squeeze(-1)
 52
 53class StandardResNet(nn.Module):
 54    """Control: same state and branches, but updates all coordinates residually."""
 55    def __init__(self,d,q=32,r=1,K=4,width=48):
 56        super().__init__(); self.K=K; self.q=q; self.r=r
 57        self.embed=nn.Sequential(nn.Linear(d,q),nn.Tanh()); self.s0=nn.Parameter(torch.zeros(r))
 58        self.branches=nn.ModuleList([nn.Sequential(nn.Linear(q+r,width),nn.Tanh(),
 59            nn.Linear(width,q+r),nn.Tanh()) for _ in range(K)])
 60        self.head=nn.Sequential(nn.Linear(q+r,width),nn.Tanh(),nn.Linear(width,1))
 61    def forward(self,x):
 62        z=torch.cat([self.embed(x),self.s0.expand(x.shape[0],-1)],1)
 63        for b in self.branches: z=z+b(z)
 64        return self.head(z).squeeze(-1)
 65
 66def target(x):
 67    d=x.shape[1]
 68    # A high-dimensional, bounded analogue of a heat/Picard target: additive
 69    # first-order term plus a smooth nonlinear correction.
 70    a=torch.arange(1,d+1,device=x.device,dtype=x.dtype)
 71    return (torch.sin(x)*((a%7+1)/4)).sum(1)/math.sqrt(d) + 0.25*torch.tanh(x.sum(1)/math.sqrt(d))
 72
 73def count(m): return sum(p.numel() for p in m.parameters() if p.requires_grad)
 74
 75def math_check():
 76    rng=np.random.default_rng(7); n=30; q=5; r=2
 77    h=rng.normal(size=(n,q)); s=rng.normal(size=(n,r)); original=s.copy(); ds=[]
 78    for k in range(6):
 79        delta=rng.normal(size=(n,r))/(k+1); ds.append(delta); s=s+delta
 80    err=float(np.max(np.abs(s-(original+sum(ds)))))
 81    # Explicit composition of two additive blocks equals one concatenated rollout.
 82    x=rng.normal(size=(n,3)); A=rng.normal(size=(3,2)); B=rng.normal(size=(2,2)); C=rng.normal(size=(2,2))
 83    y=x@A; y1=y+y@B; y2=y1+y1@C
 84    composed=(x@A)+((x@A)@B)+((x@A+(x@A)@B)@C)
 85    return {'telescoping_max_error':err,'composition_max_error':float(np.max(np.abs(y2-composed))),
 86            'bounded_tanh_max_abs':float(np.max(np.abs(np.tanh(rng.normal(size=10000))))) }
 87
 88def make_data(d, n, dev, seed):
 89    g=torch.Generator(device='cpu').manual_seed(seed)
 90    x=torch.randn(n,d,generator=g)
 91    return x.to(dev), target(x.to(dev))
 92
 93def train(model,x,y,xt,yt,steps=450,batch=128):
 94    opt=torch.optim.Adam(model.parameters(),lr=2e-3); N=len(x); hist=[]
 95    model.train(); t0=time.time()
 96    for step in range(steps):
 97        ix=torch.randint(N,(min(batch,N),),device=x.device)
 98        loss=((model(x[ix])-y[ix])**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
 99        if step in (0,99,249,449):
100            model.eval()
101            with torch.no_grad(): hist.append(float(((model(xt)-yt)**2).mean().cpu()))
102            model.train()
103    model.eval()
104    with torch.no_grad(): mse=float(((model(xt)-yt)**2).mean().cpu())
105    return mse,hist,time.time()-t0
106
107def main():
108    seed_all(); dev=device_or_cpu(); print('device',dev)
109    result={'device':str(dev),'math_check':math_check(),'runs':[]}
110    for d in (50,100,200):
111        x,y=make_data(d,1800,dev,100+d); xt,yt=make_data(d,600,dev,900+d)
112        for name, cls in [('plain_mlp',PlainMLP),('standard_resnet',StandardResNet),('accumulator',AccumulatorResNet)]:
113            seed_all(SEED+d+len(name)); m=cls(d).to(dev)
114            try: mse,hist,secs=train(m,x,y,xt,yt)
115            except Exception as e:
116                if dev.type=='cuda':
117                    print('CUDA error, rerun CPU:',repr(e)); dev=torch.device('cpu'); x,y,xt,yt=[z.cpu() for z in (x,y,xt,yt)]; m=cls(d).to(dev); mse,hist,secs=train(m,x,y,xt,yt)
118                else: raise
119            result['runs'].append({'d':d,'model':name,'params':count(m),'test_mse':mse,'checkpoints':hist,'seconds':secs})
120            print(d,name,count(m),mse,hist)
121    Path('results.json').write_text(json.dumps(result,indent=2))
122    print('wrote results.json')
123if __name__=='__main__': main()