1import json, math, random, time
 2from pathlib import Path
 3import numpy as np
 4import torch
 5import torch.nn as nn
 6import torch.nn.functional as F
 7
 8SEED=275
 9random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
10torch.set_num_threads(4)
11try:
12    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
13    if device.type=='cuda':
14        torch.zeros(1, device=device)
15except Exception:
16    device=torch.device('cpu')
17
18# ---- exact finite empirical mean-field layer ----
19class ParticleNet(nn.Module):
20    def __init__(self,d,m):
21        super().__init__(); self.m=m; self.d=d
22        self.a=nn.Parameter(0.25*torch.randn(m, device=device))
23        self.w=nn.Parameter(torch.randn(m,d, device=device)/math.sqrt(d))
24    def forward(self,x):
25        return (F.relu(x @ self.w.T)*self.a).mean(1,keepdim=True)
26    def project(self,A=3.0):
27        with torch.no_grad():
28            self.a.clamp_(-A,A)
29
30class DenseNet(nn.Module):
31    def __init__(self,d,m):
32        super().__init__(); self.w=nn.Parameter(torch.randn(m,d,device=device)/math.sqrt(d)); self.a=nn.Parameter(0.25*torch.randn(m,device=device))
33    def forward(self,x): return (F.relu(x@self.w.T)*self.a).mean(1,keepdim=True)
34
35def math_check():
36    torch.manual_seed(SEED+1)
37    d,m,n=5,17,23
38    a=torch.randn(m); w=torch.randn(m,d); x=torch.randn(n,d)
39    lhs=(F.relu(x@w.T)*a).mean(1)
40    rhs=torch.stack([(a*F.relu(w@x[i])).mean() for i in range(n)])
41    pred_err=(lhs-rhs).abs().max().item()
42    # Langevin variance: increments with zero gradient have variance 2 eta lambda.
43    eta,lam=0.013,0.7; reps=20000
44    z=torch.randn(reps)*math.sqrt(2*eta*lam)
45    var=z.var(unbiased=True).item(); target=2*eta*lam
46    return {'predictor_max_abs_error':pred_err,'noise_variance':var,'noise_target':target,'relative_variance_error':abs(var-target)/target}
47
48def make_data(d=20,ntrain=2400,ntest=1200):
49    g=torch.Generator(device=device).manual_seed(SEED+4)
50    # two distinct teacher directions: multi-index test of retained diversity
51    q,_=torch.linalg.qr(torch.randn(d,2,device=device,generator=g))
52    Xtr=torch.randn(ntrain,d,device=device,generator=g); Xte=torch.randn(ntest,d,device=device,generator=g)
53    def target(x):
54        return (0.9*F.relu(x@q[:,0])+0.7*F.relu(x@q[:,1])-0.25*F.relu(-x@q[:,0])).unsqueeze(1)
55    ytr=target(Xtr); yte=target(Xte)
56    return Xtr,ytr,Xte,yte,q
57
58def alignment(model,q):
59    # max absolute cosine to either teacher direction, and fraction of particles assigned each mode
60    wn=F.normalize(model.w,dim=1); qn=F.normalize(q,dim=0)
61    c=(wn@qn).abs(); best=c.max(1).values
62    assign=c.argmax(1)
63    return float(best.mean()), float((best>0.7).float().mean()), [int((assign==k).sum()) for k in range(2)]
64
65def train(kind,X,y,Xte,yte,q,m=64,steps=700,batch=96):
66    torch.manual_seed(SEED+9)
67    model=ParticleNet(X.shape[1],m) if kind=='particle' else DenseNet(X.shape[1],m)
68    model.to(device); opt=torch.optim.Adam(model.parameters(),lr=0.025)
69    losses=[]; t0=time.time(); lam_hi,lam_lo=1.0,0.01; anneal=steps
70    for t in range(steps):
71        ix=torch.randint(0,X.shape[0],(batch,),device=device)
72        loss=F.mse_loss(model(X[ix]),y[ix]); opt.zero_grad(); loss.backward()
73        # Adam is used for fair neural optimization; particle dynamics adds the specified diffusion.
74        opt.step()
75        if kind=='particle':
76            lam=lam_lo+0.5*(lam_hi-lam_lo)*(1+math.cos(math.pi*t/anneal))
77            eta=0.025
78            with torch.no_grad():
79                model.a.add_(math.sqrt(2*eta*lam)*torch.randn_like(model.a))
80                model.w.add_(math.sqrt(2*eta*lam)*torch.randn_like(model.w))
81                model.project()
82        if t in [0,99,299,699]:
83            with torch.no_grad(): losses.append(float(F.mse_loss(model(Xte),yte)))
84    with torch.no_grad(): test=float(F.mse_loss(model(Xte),yte)); trainloss=float(F.mse_loss(model(X),y))
85    al=alignment(model,q)
86    diversity=float(torch.pdist(F.normalize(model.w,dim=1)).mean())
87    return {'test_mse':test,'train_mse':trainloss,'checkpoints':losses,'alignment_mean':al[0],'aligned_fraction':al[1],'mode_counts':al[2],'weight_diversity':diversity,'seconds':time.time()-t0}
88
89def main():
90    check=math_check(); X,y,Xte,yte,q=make_data()
91    base=train('baseline',X,y,Xte,yte,q); idea=train('particle',X,y,Xte,yte,q)
92    out={'device':str(device),'seed':SEED,'math_check':check,'baseline':base,'particle':idea,'config':{'width':64,'dimension':20,'steps':700,'batch':96,'schedule':'lambda 1.0 -> 0.01 cosine','teacher':'two-index ReLU regression'}}
93    Path('results.json').write_text(json.dumps(out,indent=2))
94    print(json.dumps(out,indent=2))
95if __name__=='__main__': main()