1import json, math, random
  2import numpy as np
  3import torch
  4
  5SEED = 2235
  6np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  7torch.set_num_threads(min(12, torch.get_num_threads()))
  8try:
  9    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 10    if device.type == 'cuda':
 11        torch.cuda.set_device(0)
 12        torch.empty(1, device=device)
 13except Exception:
 14    device = torch.device('cpu')
 15
 16# Smooth exchangeable module: f(W)=sum_i v_i phi(u_i^T x), phi(0)=0.
 17# The quadratic term is phi'(0) sum_i v_i u_i^T x.
 18def phi(s):
 19    return s + 0.5*s*s + 0.1*s*s*s
 20
 21def branch_output(W, x):
 22    # W is p x d, last row is v and first k rows are u.
 23    u, v = W[:-1, :], W[-1, :]
 24    return (v * phi(u.T @ x)).sum()
 25
 26def structure_A(x):
 27    k = x.numel()
 28    A = torch.zeros((k+1, k+1), dtype=x.dtype, device=x.device)
 29    A[:-1, -1] = 0.5*x
 30    A[-1, :-1] = 0.5*x
 31    return A
 32
 33def quadratic_output(W, x):
 34    M = W @ W.T
 35    return torch.sum(M * structure_A(x))
 36
 37def fit_slope(xs, ys):
 38    return float(np.polyfit(np.log(xs), np.log(np.maximum(ys, 1e-30)), 1)[0])
 39
 40def math_checks():
 41    k, d = 5, 13
 42    x = torch.randn(k)
 43    base = torch.randn(k+1, d)
 44    eps = np.logspace(-3, -0.5, 8)
 45    residuals=[]; magnitudes=[]
 46    for e in eps:
 47        W=e*base
 48        residuals.append(abs(float(branch_output(W,x)-quadratic_output(W,x))))
 49        magnitudes.append(float(torch.linalg.norm(W)))
 50    slope=fit_slope(np.array(magnitudes), np.array(residuals))
 51
 52    # Exact exchangeability: arbitrary branch permutation leaves both outputs unchanged.
 53    perm=torch.randperm(d)
 54    perm_err=abs(float(branch_output(base[:,perm],x)-branch_output(base,x)))
 55    # Gram-only quadratic invariance: orthogonal mixing leaves WW^T exactly unchanged,
 56    # while the original module differs only by higher-order terms near zero.
 57    Q,_=torch.linalg.qr(torch.randn(d,d))
 58    gram_err=abs(float(quadratic_output(base@Q,x)-quadratic_output(base,x)))
 59    ortho_errs=[]
 60    for e in eps:
 61        ortho_errs.append(abs(float(branch_output(e*base@Q,x)-branch_output(e*base,x))))
 62    ortho_slope=fit_slope(np.array(magnitudes), np.array(ortho_errs))
 63
 64    # Width sweep: independent branch cubic contributions predict RMS residual ~ d^1/2
 65    # for random signs, while a fixed aligned construction predicts linear d. Measure both.
 66    widths=np.array([4,8,16,32,64])
 67    rms=[]; aligned=[]
 68    # Use nested prefixes for random widths and one fixed column/input for the
 69    # coherent construction; otherwise resampling would add unrelated noise.
 70    fixed_x=torch.randn(k)
 71    fixed_col=torch.randn(k+1,1)
 72    for dd in widths:
 73        vals=[]
 74        for rep in range(30):
 75            X=torch.randn(k+1,dd); xx=torch.randn(k)
 76            vals.append(abs(float(branch_output(0.03*X,xx)-quadratic_output(0.03*X,xx))))
 77        rms.append(np.mean(vals))
 78        X=fixed_col.repeat(1,dd)
 79        aligned.append(abs(float(branch_output(0.03*X,fixed_x)-quadratic_output(0.03*X,fixed_x))))
 80    width_slope_rms=fit_slope(widths, np.array(rms))
 81    width_slope_aligned=fit_slope(widths, np.array(aligned))
 82    return {
 83      'device':str(device), 'epsilon_values':eps.tolist(),
 84      'remainder_values':residuals, 'remainder_loglog_slope':slope,
 85      'predicted_remainder_slope':3.0,
 86      'permutation_error':perm_err, 'quadratic_orthogonal_error':gram_err,
 87      'original_orthogonal_difference_slope':ortho_slope,
 88      'predicted_orthogonal_difference_slope':3.0,
 89      'widths':widths.tolist(), 'width_mean_abs_remainder':rms,
 90      'width_slope_random':width_slope_rms, 'predicted_width_slope_random':0.5,
 91      'width_slope_aligned':width_slope_aligned, 'predicted_width_slope_aligned':1.0
 92    }
 93
 94class Branch(torch.nn.Module):
 95    def __init__(self,k,d,init=0.05):
 96        super().__init__(); self.W=torch.nn.Parameter(init*torch.randn(k+1,d))
 97    def forward(self,x):
 98        u,v=self.W[:-1],self.W[-1]
 99        return (phi(x @ u) * v).sum(1)
100class Quadratic(torch.nn.Module):
101    def __init__(self,k,d,init=0.05):
102        super().__init__(); self.W=torch.nn.Parameter(init*torch.randn(k+1,d))
103    def forward(self,x):
104        # Equivalent low-cost factorized trace: sum_i w_i^T A(x) w_i.
105        u,v=self.W[:-1],self.W[-1]
106        return (x @ u * v).sum(1)
107
108def train_model(cls, X, y, k, d, steps=500):
109    torch.manual_seed(SEED+7)
110    m=cls(k,d).to(device); opt=torch.optim.Adam(m.parameters(),lr=0.03)
111    Xt,yt=X.to(device),y.to(device)
112    losses=[]
113    for t in range(steps):
114        opt.zero_grad(); pred=m(Xt); loss=((pred-yt)**2).mean(); loss.backward(); opt.step()
115        if t in (0,99,249,499): losses.append(float(loss.detach().cpu()))
116    return losses
117
118def mini_experiment():
119    torch.manual_seed(SEED+99)
120    k,d,N=6,16,128
121    X=torch.randn(N,k)
122    teacher=torch.randn(k+1,d)*0.18
123    y=branch_output(teacher,X.T if False else X[0]) # overwritten vectorized below
124    u,v=teacher[:-1],teacher[-1]
125    y=(phi(X@u)*v).sum(1)
126    b=train_model(Branch,X,y,k,d)
127    q=train_model(Quadratic,X,y,k,d)
128    return {'dataset':'synthetic cubic exchangeable branch teacher','k':k,'width':d,
129            'steps':[0,99,249,499],'baseline_branch_mse':b,'idea_quadratic_mse':q,
130            'parameter_count_each':(k+1)*d}
131
132def main():
133    checks=math_checks(); comparison=mini_experiment()
134    # Mechanism passes only if all requested scaling predictions are numerically close.
135    checks['prediction_pass'] = (abs(checks['remainder_loglog_slope']-3)<0.15 and
136        abs(checks['original_orthogonal_difference_slope']-3)<0.2 and
137        abs(checks['width_slope_random']-0.5)<0.25 and
138        abs(checks['width_slope_aligned']-1)<0.2 and
139        checks['quadratic_orthogonal_error']<1e-5 and checks['permutation_error']<1e-5)
140    out={'math_checks':checks,'mini_experiment':comparison}
141    with open('results.json','w') as f: json.dump(out,f,indent=2)
142    print(json.dumps(out,indent=2))
143if __name__=='__main__': main()