1import json, math, random
  2import numpy as np
  3import torch
  4from torch import nn
  5import sys
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, train_model, evaluate, sweep_baseline, make_report
  8
  9SEED = 2235
 10WIDTH = 64
 11HIDDEN = 16
 12NTR = 800
 13NTE = 400
 14SEEDS = tuple(range(8))
 15# Smooth activation used by the standard branch MLP; its first two Taylor
 16# terms produce the stated quadratic Gram module.
 17def phi(s):
 18    return s + 0.5*s*s + 0.1*s*s*s
 19
 20def math_check():
 21    torch.manual_seed(SEED)
 22    n, d, h = 7, WIDTH, HIDDEN
 23    x = torch.randn(19, n, dtype=torch.float64)
 24    base = torch.randn(n+h, d, dtype=torch.float64)
 25    def full(W):
 26        u, v = W[:n], W[n:]
 27        return (phi(x @ u) @ v.T)
 28    def quad(W):
 29        u, v = W[:n], W[n:]
 30        # Tr(WW^T A_j(x)) = sum_i v_ij (x^T u_i)
 31        return ((x @ u) @ v.T)
 32    es = np.logspace(-3, -0.7, 8)
 33    rem = np.array([float((full(e*base)-quad(e*base)).abs().mean()) for e in es])
 34    slope = float(np.polyfit(np.log(es), np.log(rem), 1)[0])
 35    q, _ = torch.linalg.qr(torch.randn(d,d, dtype=torch.float64))
 36    gram_err = float((quad(base@q)-quad(base)).abs().max())
 37    return {'epsilon_values': es.tolist(), 'remainder_values': rem.tolist(),
 38            'observed_remainder_slope': slope, 'predicted_remainder_slope': 3.0,
 39            'gram_invariance_max_error': gram_err,
 40            'prediction_pass': bool(abs(slope-3.0)<0.15 and gram_err<1e-5)}
 41
 42class BranchMLP(nn.Module):
 43    def __init__(self, n, idea=False):
 44        super().__init__(); self.idea=idea
 45        self.W=nn.Parameter(0.05*torch.randn(n+HIDDEN, WIDTH))
 46        self.proj=nn.Linear(HIDDEN, 1)
 47    def forward(self, x):
 48        u,v=self.W[:x.shape[-1]],self.W[x.shape[-1]:]
 49        a=x @ u
 50        z = (a @ v.T) if self.idea else (phi(a) @ v.T)
 51        return self.proj(z)
 52
 53def make_train_fn(cfg, idea):
 54    def run(seed):
 55        torch.manual_seed(SEED + 1000 + int(seed)); np.random.seed(SEED+int(seed)); random.seed(SEED+int(seed))
 56        d=get_dataset('tabular', int(seed), n_train=NTR, n_test=NTE)
 57        m=BranchMLP(d['input_shape'][0], idea=idea)
 58        _, metric, _ = train_model(m, d, epochs=30, lr=float(cfg['lr']), batch=128, weight_decay=float(cfg['weight_decay']), log=lambda *_: None)
 59        return float(metric)
 60    return run
 61
 62def trained_signature():
 63    # Obtain a trained baseline model and re-test the cubic remainder prediction
 64    # on its actual learned branch weights, rather than on a synthetic graph.
 65    seed=0; d=get_dataset('tabular', seed, n_train=NTR, n_test=NTE)
 66    torch.manual_seed(SEED+1000); m=BranchMLP(d['input_shape'][0], idea=False)
 67    m,_,_=train_model(m,d,epochs=30,lr=0.003,batch=128,weight_decay=0.0,log=lambda *_: None)
 68    m = m.cpu(); x=d['xte'][:128].cpu(); W=m.W.detach().cpu(); n=x.shape[1]; u,v=W[:n],W[n:]
 69    def f(ww):
 70        uu,vv=ww[:n],ww[n:]; return (phi(x@uu)@vv.T) @ m.proj.weight.detach().T + m.proj.bias.detach()
 71    def q(ww):
 72        uu,vv=ww[:n],ww[n:]; return ((x@uu)@vv.T) @ m.proj.weight.detach().T + m.proj.bias.detach()
 73    es=np.logspace(-2.5,-0.8,7); vals=np.array([float((f(e*W)-q(e*W)).abs().mean()) for e in es])
 74    slope=float(np.polyfit(np.log(es),np.log(np.maximum(vals,1e-30)),1)[0])
 75    perm=torch.randperm(W.shape[1]); perr=float((f(W[:,perm])-f(W)).abs().max())
 76    return {'source':'trained baseline model on tabular test inputs','epsilon_values':es.tolist(),
 77            'observed_remainder_slope':slope,'predicted_remainder_slope':3.0,
 78            'permutation_max_error':perr,
 79            'confirmed':bool(abs(slope-3.0)<0.25 and perr<1e-5)}
 80
 81def main():
 82    checks=math_check(); print('MATH',json.dumps(checks))
 83    grid=[{'lr':0.001,'weight_decay':0.0},{'lr':0.003,'weight_decay':0.0},{'lr':0.01,'weight_decay':0.0}]
 84    base=sweep_baseline(lambda c: make_train_fn(c,False), grid, seeds=(0,1,2,3))
 85    idea_trials=[]
 86    for cfg in grid:
 87        r=evaluate(make_train_fn(cfg,True), SEEDS)
 88        idea_trials.append({'cfg':cfg,'result':r})
 89    best=min(idea_trials,key=lambda z:z['result']['mean'])
 90    idea={'best_cfg':best['cfg'],'sweep':[{'cfg':z['cfg'],'mean':z['result']['mean']} for z in idea_trials],
 91          'per_seed':best['result']['per_seed'],'mean':best['result']['mean'],'std':best['result']['std'],'n':best['result']['n']}
 92    sig=trained_signature()
 93    rep=make_report('tabular','custom_branch_mlp',base,idea,{'math_check':checks,'trained_model':sig,
 94        'custom_track':None})
 95    rep['protocol_notes']={'paired_seeds':list(SEEDS),'n_train':NTR,'n_test':NTE,
 96        'architecture':'same width-64 exchangeable branch MLP; only phi vs quadratic Taylor block differs',
 97        'baseline_grid':grid,'idea_grid':grid}
 98    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
 99    print(json.dumps(rep,indent=2))
100if __name__=='__main__': main()