1import sys, json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, make_model, sweep_baseline, make_report
  8
  9OUT = Path('bench_report.json')
 10SEEDS = tuple(range(8))
 11# Shared union: baseline is evaluated at every lr used by the idea.
 12GRID = [
 13    {'lr': 0.001, 'weight_decay': 0.0},
 14    {'lr': 0.003, 'weight_decay': 0.0},
 15    {'lr': 0.006, 'weight_decay': 0.0},
 16]
 17EPOCHS = 18
 18BATCH = 128
 19
 20
 21def seed_all(seed):
 22    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 23    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 24
 25
 26def c_transition(k):
 27    k = max(float(k), 1.0)
 28    return math.sqrt(((k+1)+math.sqrt(max(0., (k+1)**2-4)))/2)
 29
 30
 31def calibrated_t(k, delta):
 32    k = max(float(k), 1.0)
 33    raw = math.sqrt(1 + math.sqrt(max(0., (k-1)*(1/delta-1))))
 34    return max(raw, c_transition(k))
 35
 36
 37def math_check():
 38    max_inv = 0.; max_poly = 0.; min_margin = 1e9
 39    for k in [1.01, 1.1, 2., 3., 10., 100.]:
 40        c = c_transition(k)
 41        max_poly = max(max_poly, abs(c**4-(k+1)*c**2+1))
 42        for d in [.001, .01, .05, .2]:
 43            raw = math.sqrt(1 + math.sqrt((k-1)*(1/d-1)))
 44            if raw >= c:
 45                tail = (k-1)/((raw*raw-1)**2+k-1)
 46                max_inv = max(max_inv, abs(tail-d))
 47            min_margin = min(min_margin, calibrated_t(k,d)-c)
 48    return {'max_inverse_error': max_inv, 'max_transition_polynomial_error': max_poly,
 49            'min_projected_tail_margin': min_margin, 'passed': max_inv < 1e-12 and max_poly < 1e-10}
 50
 51
 52def run(seed, cfg, idea=False, collect=False):
 53    seed_all(seed)
 54    ds = get_dataset('tabular', seed, n_train=400, n_test=400)
 55    # Explicitly use the matched canonical architecture.
 56    net = make_model('mlp_tiny', ds['input_shape'], ds['out_dim'])
 57    device = 'cuda' if torch.cuda.is_available() else 'cpu'
 58    try:
 59        net = net.to(device)
 60    except Exception:
 61        device = 'cpu'; net = net.to(device)
 62    opt = torch.optim.Adam(net.parameters(), lr=cfg['lr'], weight_decay=cfg.get('weight_decay',0.))
 63    x, y = ds['xtr'].to(device), ds['ytr'].to(device)
 64    lossf = nn.MSELoss(); n = len(x)
 65    beta=.95; delta=.02; safety=1.5; cap=1000.
 66    stats={'raw_spike':0, 'clipped_entries':0, 'entries':0, 'tail_sum':0., 'k_sum':0., 'tau_sum':0., 'steps':0}
 67    for ep in range(EPOCHS):
 68        perm = torch.randperm(n, device=device)
 69        for start in range(0,n,BATCH):
 70            ix=perm[start:start+BATCH]; opt.zero_grad(set_to_none=True)
 71            loss=lossf(net(x[ix]), y[ix]); loss.backward()
 72            for p in net.parameters():
 73                if p.grad is None: continue
 74                g=p.grad
 75                if not idea:
 76                    stats['raw_spike'] += int(torch.isfinite(g).all() and g.detach().abs().max().item() > 5.0)
 77                    continue
 78                z=g.detach(); m=z.mean(); centered=z-m
 79                v_batch=(centered*centered).mean(); q_batch=(centered**4).mean()
 80                # Per-tensor EMA moments; current-batch central moments avoid lag bias.
 81                if not hasattr(p, '_kc_state'):
 82                    p._kc_state=[m.item(), max(v_batch.item(),1e-12), max(q_batch.item(),1e-12)]
 83                st=p._kc_state
 84                st[0]=beta*st[0]+(1-beta)*m.item(); st[1]=beta*st[1]+(1-beta)*v_batch.item(); st[2]=beta*st[2]+(1-beta)*q_batch.item()
 85                k=min(cap,max(1., safety*st[2]/(st[1]**2+1e-12))); t=calibrated_t(k,delta); tau=t*math.sqrt(st[1]+1e-12)
 86                stats['raw_spike'] += int(torch.isfinite(g).all() and z.abs().max().item()>5.0)
 87                stats['clipped_entries'] += int((centered.abs()>tau).sum().item()); stats['entries'] += z.numel()
 88                stats['tail_sum'] += float((centered.abs()>tau).float().mean().item()); stats['k_sum'] += k; stats['tau_sum'] += tau
 89                g.copy_(m + centered.clamp(-tau,tau))
 90            opt.step(); stats['steps'] += 1
 91    with torch.no_grad(): metric=float(lossf(net(ds['xte'].to(device)), ds['yte'].to(device)).item())
 92    if collect:
 93        stats['mean_tail_above_tau']=stats['tail_sum']/max(stats['steps']*len(list(net.parameters())),1)
 94        stats['mean_kurtosis_bound']=stats['k_sum']/max(stats['steps']*len(list(net.parameters())),1)
 95        stats['mean_tau']=stats['tau_sum']/max(stats['steps']*len(list(net.parameters())),1)
 96        stats['clip_fraction']=stats['clipped_entries']/max(stats['entries'],1)
 97        return metric, stats
 98    return metric
 99
100
101def fn(idea, cfg): return lambda s: run(s,cfg,idea=idea)
102
103# Baseline sweep uses exactly the idea's three learning rates and four sweep seeds.
104base = sweep_baseline(lambda cfg: fn(False,cfg), GRID, seeds=(0,1,2,3))
105best = base['best_cfg']
106# Idea at best baseline lr and two nearby settings, same grid; choose by the same 4-seed tuning protocol.
107idea_trials=[]
108for cfg in GRID:
109    r=[run(s,cfg,idea=True) for s in (0,1,2,3)]
110    idea_trials.append({'cfg':cfg,'mean':float(np.mean(r))})
111idea_cfg=min(GRID,key=lambda c: next(q['mean'] for q in idea_trials if q['cfg']==c))
112idea_res={'cfg':idea_cfg, **__import__('bench').protocol.evaluate(fn(True,idea_cfg), seeds=SEEDS)}
113# Collect trained-model behavior on all paired seeds (not an analytical toy signature).
114beh=[]
115for s in SEEDS:
116    val, st=run(s,idea_cfg,idea=True,collect=True); beh.append(st)
117mean_tail=float(np.mean([z['mean_tail_above_tau'] for z in beh])); mean_clip=float(np.mean([z['clip_fraction'] for z in beh]))
118mean_k=float(np.mean([z['mean_kurtosis_bound'] for z in beh]))
119# Prediction: calibrated threshold should make observed two-sided exceedance no larger than target order.
120signature={'prediction':'trained gradient-coordinate tail above calibrated tau is controlled near target delta=.02',
121           'target_delta':0.02, 'observed_mean_tail_above_tau':mean_tail,
122           'mean_clip_fraction':mean_clip, 'mean_kurtosis_bound':mean_k,
123           'confirmed': bool(mean_tail <= 0.02*2.0)}
124report=make_report('tabular','mlp_tiny',base,idea_res,{'mechanism_signature':signature,
125    'math_check':math_check(),'idea_sweep':idea_trials,
126    'protocol_note':'8 paired seeds; baseline sweep and idea sweep share the full learning-rate union; Adam and architecture are otherwise identical.'})
127report['idea']['behavior_all_seeds']=beh
128OUT.write_text(json.dumps(report,indent=2))
129print(json.dumps(report,indent=2))