import sys, json, math, random from pathlib import Path import numpy as np import torch import torch.nn as nn sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import get_dataset, train_model, sweep_baseline, make_report from bench.protocol import evaluate SEEDS = tuple(range(8)) SWEEP_SEEDS = (0, 1, 2, 3) EPOCHS = 18 NTRAIN, NTEST = 2000, 500 D, DEPTH, HEADS = 64, 2, 2 S, SIGMA, LAMBDA = 0.8, 0.35, 0.7 def kernel_np(x, y, d=3, s=S, sigma=SIGMA, eps=1e-5): r = np.linalg.norm(x-y, axis=-1) + eps ux = np.linalg.norm(x, axis=-1) + eps uy = np.linalg.norm(y, axis=-1) + eps m = np.minimum(np.minimum(ux/r, uy/r), 1.0) return r**(s-d) * m**(-sigma) def math_check(): rng = np.random.RandomState(7) x, y = rng.uniform(-1, 1, (2, 20, 3)) k1 = kernel_np(x, y) k2 = kernel_np(3*x, 3*y) hom = np.max(np.abs(k2/(k1*3**(S-3))-1)) near = np.array([[1e-5, 0, 0.]]) far = np.array([[.8, 0, 0.]]) amp = float(kernel_np(near, far)[0] / ((np.linalg.norm(near-far)+1e-5)**(S-3))) pred = float(((np.linalg.norm(near[0])+1e-5)/(np.linalg.norm(near[0]-far[0])+1e-5))**(-SIGMA)) return {'homogeneity_max_relative_error': float(hom), 'near_origin_amplification': amp, 'predicted_amplification': pred, 'admissible_d3': bool(0 < S < 3-2*SIGMA), 'relative_amplification_error': float(abs(amp/pred-1))} def coords3(length, device): # A bounded 3-D temporal embedding; the forecast endpoint is the singular center. t = torch.linspace(-1., 0., length, device=device) return torch.stack((t, .35*torch.sin(math.pi*t), .35*torch.cos(math.pi*t)-.35), -1) class CoordTransformer(nn.Module): def __init__(self, win, out_dim, biased=False): super().__init__(); self.win=win; self.biased=biased self.inp=nn.Linear(1, D); self.pos=nn.Parameter(torch.zeros(1, win, D)) nn.init.normal_(self.pos, std=.02) self.layers=nn.ModuleList() for _ in range(DEPTH): self.layers.append(nn.ModuleDict({ 'norm1': nn.LayerNorm(D), 'attn': nn.MultiheadAttention(D, HEADS, dropout=0., batch_first=True), 'norm2': nn.LayerNorm(D), 'ff': nn.Sequential(nn.Linear(D,128), nn.ReLU(), nn.Linear(128,D))})) self.head=nn.Linear(win*D, out_dim) self.last_attn=None def forward(self, x): b, n = x.shape h=self.inp(x.unsqueeze(-1))+self.pos[:, :n] c=coords3(n, x.device) bias=None if self.biased: dif=c[:,None,:]-c[None,:,:] r=torch.sqrt((dif*dif).sum(-1)+1e-10)+1e-5 u=torch.sqrt((c*c).sum(-1)+1e-10)+1e-5 m=torch.minimum(torch.minimum(u[:,None]/r,u[None,:]/r),torch.ones_like(r)) K=r**(S-3)*m**(-SIGMA) bias=(LAMBDA*torch.log(K+1e-8)).unsqueeze(0).expand(b,-1,-1) bias=bias.repeat_interleave(HEADS, 0) weights=None for li in self.layers: z=li['norm1'](h) h2, weights=li['attn'](z,z,z, attn_mask=bias, need_weights=True, average_attn_weights=False) h=h+h2; h=h+li['ff'](li['norm2'](h)) self.last_attn=weights.detach() if weights is not None else None return self.head(h.reshape(b,-1)) def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) def train_one(cfg, seed, biased, capture=False): seed_all(seed) ds=get_dataset('sequence', seed, n_train=NTRAIN, n_test=NTEST) net=CoordTransformer(ds['input_shape'][0], ds['out_dim'], biased=biased) trained, metric, hist=train_model(net, ds, epochs=EPOCHS, lr=cfg['lr'], batch=128, weight_decay=cfg['weight_decay'], log=lambda *_: None) if trained is None: raise RuntimeError('bench training failed') if capture: trained.eval() with torch.no_grad(): dev=next(trained.parameters()).device _=trained(ds['xte'][:64].to(dev)) att=trained.last_attn.cpu().numpy() # [B,H,N,N] return float(metric), {'model': trained, 'ds': ds, 'att': att} return float(metric) def make_train(cfg, biased): return lambda seed: train_one(cfg, seed, biased) def main(): checks=math_check() # Three learning rates are the shared union of both searches; baseline has no extra method knob. grid=[{'lr':lr, 'weight_decay':wd} for lr,wd in [(0.0015,0.0),(0.003,0.0),(0.006,0.0)]] base=sweep_baseline(lambda cfg: make_train(cfg, False), grid, seeds=SWEEP_SEEDS) # Same 3-point budget and same lr union for the idea; lambda is fixed a priori from stage 1. idea_cfgs=grid idea_sweep=[] for cfg in idea_cfgs: r=evaluate(make_train(cfg, True), seeds=SWEEP_SEEDS) idea_sweep.append({'cfg':cfg, 'mean':r['mean']}) best_idea_cfg=min(idea_cfgs, key=lambda c: next(z['mean'] for z in idea_sweep if z['cfg']==c)) idea_full=evaluate(make_train(best_idea_cfg, True), seeds=SEEDS) # Re-test trained-model attention behavior on paired seed 0, not an analytic-only signature. bm, bo=train_one(base['best_cfg'], 0, False, True) im, io=train_one(best_idea_cfg, 0, True, True) a0=bo['att']; a1=io['att']; n=a0.shape[-1] # Query endpoint (last token): compare attention odds on oldest vs latest key. old, recent = float(a0[:,:,-1,0].mean()), float(a0[:,:,-1,-1].mean()) old_i, recent_i = float(a1[:,:,-1,0].mean()), float(a1[:,:,-1,-1].mean()) observed_ratio=(recent_i/(old_i+1e-12))/(recent/(old+1e-12)) c=np.stack([np.linspace(-1,0,n), .35*np.sin(np.pi*np.linspace(-1,0,n)), .35*np.cos(np.pi*np.linspace(-1,0,n))-.35],-1) predicted_ratio=float(kernel_np(c[-1:], c[-1:])[0]/kernel_np(c[-1:], c[:1])[0]) signature={'predicted_near_vs_far_kernel_ratio':predicted_ratio, 'observed_attention_odds_ratio_after_bias':observed_ratio, 'baseline_endpoint_recent_attention':recent, 'idea_endpoint_recent_attention':recent_i, 'confirmed': bool(observed_ratio > 1.0 and observed_ratio/predicted_ratio > 1/3 and observed_ratio/predicted_ratio < 3)} report=make_report('sequence','transformer_tiny', {'best_cfg':base['best_cfg'],'sweep':base['sweep'],'full':base['full']}, idea_full, {'mechanism_signature':signature,'idea_sweep':idea_sweep,'math_check':checks, 'idea_best_cfg':best_idea_cfg,'budget':{'epochs':EPOCHS,'n_train':NTRAIN,'n_test':NTEST,'seeds':list(SEEDS)}}) Path('bench_report.json').write_text(json.dumps(report, indent=2)) print(json.dumps(report, indent=2)) if __name__=='__main__': main()