1import json, math, os, time
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5
  6SEED = 559
  7np.random.seed(SEED); torch.manual_seed(SEED)
  8try:
  9    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 10except Exception:
 11    device = torch.device("cpu")
 12
 13
 14def forward(x, w, b, a, v, c):
 15    return torch.relu(x @ w.T + b) @ a + x @ v + c
 16
 17
 18def canonical_merge(w, b, a, tol=1e-7):
 19    """Exact/tolerance merge, preserving orientation (u,beta), never sign-flips."""
 20    w = np.asarray(w, dtype=np.float64); b=np.asarray(b, dtype=np.float64); a=np.asarray(a, dtype=np.float64)
 21    groups=[]; zero=[]
 22    for j in range(len(a)):
 23        r=np.linalg.norm(w[j])
 24        if r <= tol:
 25            zero.append(j); continue
 26        u=w[j]/r; beta=b[j]/r; alpha=a[j]*r
 27        found=None
 28        for g in groups:
 29            if np.max(np.abs(np.r_[u,beta]-np.r_[g['u'],g['beta']])) <= tol:
 30                found=g; break
 31        if found is None:
 32            groups.append({'u':u.copy(),'beta':float(beta),'alpha':float(alpha)})
 33        else: found['alpha'] += float(alpha)
 34    if not groups:
 35        return np.zeros((0,w.shape[1])), np.zeros(0), np.zeros(0), zero
 36    W=np.array([g['u'] for g in groups]); B=np.array([g['beta'] for g in groups]); A=np.array([g['alpha'] for g in groups])
 37    return W,B,A,zero
 38
 39
 40def jacobian_norm(model_tuple, x):
 41    x=x.detach().clone().requires_grad_(True)
 42    y=forward(x,*model_tuple)
 43    grads=torch.autograd.grad(y.sum(),x)[0]
 44    return float(torch.linalg.vector_norm(grads,dim=1).mean())
 45
 46
 47def main():
 48    d=2; n=600
 49    x=torch.randn(n,d,device=device)
 50    y=(torch.sin(1.7*x[:,0])+0.35*x[:,1]**2-0.2*x[:,0]*x[:,1]+0.08*torch.randn(n,device=device))
 51    xv=torch.randn(400,d,device=device)
 52    yv=(torch.sin(1.7*xv[:,0])+0.35*xv[:,1]**2-0.2*xv[:,0]*xv[:,1])
 53    # Train a small ordinary network; this is the compact representation to be duplicated.
 54    m=24
 55    w=nn.Parameter(0.8*torch.randn(m,d,device=device)); b=nn.Parameter(torch.zeros(m,device=device))
 56    a=nn.Parameter(0.3*torch.randn(m,device=device)); v=nn.Parameter(torch.zeros(d,device=device)); c=nn.Parameter(torch.zeros((),device=device))
 57    opt=torch.optim.Adam([w,b,a,v,c],lr=0.025)
 58    t0=time.time()
 59    for step in range(900):
 60        opt.zero_grad(); pred=forward(x,w,b,a,v,c); loss=((pred-y)**2).mean(); loss.backward(); opt.step()
 61    train_sec=time.time()-t0
 62    base=tuple(q.detach().clone() for q in (w,b,a,v,c))
 63    # Make a width-48 nominal model by splitting every coefficient into two identical neurons.
 64    wb,bb,ab,vb,cb=base
 65    wr=torch.repeat_interleave(wb,2,dim=0); br=torch.repeat_interleave(bb,2); ar=torch.repeat_interleave(ab/2,2)
 66    dense=(wr,br,ar,vb,cb)
 67    # Merge in CPU canonical coordinates, then evaluate the resulting model.
 68    W,B,A,zero=canonical_merge(wr.cpu().numpy(),br.cpu().numpy(),ar.cpu().numpy(),tol=1e-9)
 69    merged=tuple([torch.tensor(z,dtype=torch.float32,device=device) for z in (W,B,A)])+(vb,cb)
 70    with torch.no_grad():
 71        probe=torch.randn(5000,d,device=device)
 72        yd=forward(probe,*dense); ym=forward(probe,*merged); ycompact=forward(probe,*base)
 73        maxerr=float((yd-ym).abs().max()); rmserr=float(torch.sqrt(((yd-ym)**2).mean()))
 74        val_dense=float(((forward(xv,*dense)-yv)**2).mean()); val_merge=float(((forward(xv,*merged)-yv)**2).mean())
 75        val_comp=float(((forward(xv,*base)-yv)**2).mean())
 76    # Independent algebra check over random neurons, including non-unit scales.
 77    ww=np.random.randn(17,d); bb=np.random.randn(17); aa=np.random.randn(17)
 78    W2,B2,A2,_=canonical_merge(ww,bb,aa,tol=1e-12)
 79    z=torch.randn(1000,d); lhs=torch.relu(z@torch.tensor(ww,dtype=torch.float32).T+torch.tensor(bb,dtype=torch.float32))@torch.tensor(aa,dtype=torch.float32)
 80    rhs=torch.relu(z@torch.tensor(W2,dtype=torch.float32).T+torch.tensor(B2,dtype=torch.float32))@torch.tensor(A2,dtype=torch.float32)
 81    algebra_err=float((lhs-rhs).abs().max())
 82    jac_dense=jacobian_norm(dense,probe[:512]); jac_merge=jacobian_norm(merged,probe[:512])
 83    raw_params=2*m + m + d + 1 # width 24 compact
 84    redundant_params=2*(2*m)+(2*m)+d+1
 85    merged_params=2*m+m+d+1
 86    result={
 87      'device':str(device), 'seed':SEED, 'train_seconds':train_sec,
 88      'math_algebra_max_error':algebra_err,
 89      'constructed_width':2*m, 'merged_width':len(A), 'width_reduction_factor':(2*m)/len(A),
 90      'probe_max_abs_function_error':maxerr, 'probe_rmse_function_error':rmserr,
 91      'val_mse_redundant':val_dense, 'val_mse_merged':val_merge, 'val_mse_compact':val_comp,
 92      'jacobian_mean_norm_redundant':jac_dense, 'jacobian_mean_norm_merged':jac_merge,
 93      'params_redundant':redundant_params, 'params_merged':merged_params,
 94      'note':'The redundant model is an exact coefficient-split reparameterization of a trained compact model; this isolates the claimed merger phenomenon rather than claiming spontaneous duplicate discovery.'
 95    }
 96    os.makedirs('results',exist_ok=True)
 97    with open('results/metrics.json','w') as f: json.dump(result,f,indent=2)
 98    print(json.dumps(result,indent=2))
 99
100if __name__=='__main__':
101    try: main()
102    except Exception as e:
103        if torch.cuda.is_available() and str(device)=='cuda':
104            print('CUDA failed; rerun on CPU:',repr(e)); device=torch.device('cpu'); main()
105        else: raise