1import sys, json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import get_dataset, train_model, sweep_baseline, make_report
  9from bench.protocol import evaluate
 10
 11SEEDS = tuple(range(8))
 12SWEEP_SEEDS = (0, 1, 2, 3)
 13EPOCHS = 18
 14NTRAIN, NTEST = 2000, 500
 15D, DEPTH, HEADS = 64, 2, 2
 16S, SIGMA, LAMBDA = 0.8, 0.35, 0.7
 17
 18
 19def kernel_np(x, y, d=3, s=S, sigma=SIGMA, eps=1e-5):
 20    r = np.linalg.norm(x-y, axis=-1) + eps
 21    ux = np.linalg.norm(x, axis=-1) + eps
 22    uy = np.linalg.norm(y, axis=-1) + eps
 23    m = np.minimum(np.minimum(ux/r, uy/r), 1.0)
 24    return r**(s-d) * m**(-sigma)
 25
 26
 27def math_check():
 28    rng = np.random.RandomState(7)
 29    x, y = rng.uniform(-1, 1, (2, 20, 3))
 30    k1 = kernel_np(x, y)
 31    k2 = kernel_np(3*x, 3*y)
 32    hom = np.max(np.abs(k2/(k1*3**(S-3))-1))
 33    near = np.array([[1e-5, 0, 0.]])
 34    far = np.array([[.8, 0, 0.]])
 35    amp = float(kernel_np(near, far)[0] / ((np.linalg.norm(near-far)+1e-5)**(S-3)))
 36    pred = float(((np.linalg.norm(near[0])+1e-5)/(np.linalg.norm(near[0]-far[0])+1e-5))**(-SIGMA))
 37    return {'homogeneity_max_relative_error': float(hom),
 38            'near_origin_amplification': amp, 'predicted_amplification': pred,
 39            'admissible_d3': bool(0 < S < 3-2*SIGMA),
 40            'relative_amplification_error': float(abs(amp/pred-1))}
 41
 42
 43def coords3(length, device):
 44    # A bounded 3-D temporal embedding; the forecast endpoint is the singular center.
 45    t = torch.linspace(-1., 0., length, device=device)
 46    return torch.stack((t, .35*torch.sin(math.pi*t), .35*torch.cos(math.pi*t)-.35), -1)
 47
 48
 49class CoordTransformer(nn.Module):
 50    def __init__(self, win, out_dim, biased=False):
 51        super().__init__(); self.win=win; self.biased=biased
 52        self.inp=nn.Linear(1, D); self.pos=nn.Parameter(torch.zeros(1, win, D))
 53        nn.init.normal_(self.pos, std=.02)
 54        self.layers=nn.ModuleList()
 55        for _ in range(DEPTH):
 56            self.layers.append(nn.ModuleDict({
 57                'norm1': nn.LayerNorm(D), 'attn': nn.MultiheadAttention(D, HEADS, dropout=0., batch_first=True),
 58                'norm2': nn.LayerNorm(D), 'ff': nn.Sequential(nn.Linear(D,128), nn.ReLU(), nn.Linear(128,D))}))
 59        self.head=nn.Linear(win*D, out_dim)
 60        self.last_attn=None
 61
 62    def forward(self, x):
 63        b, n = x.shape
 64        h=self.inp(x.unsqueeze(-1))+self.pos[:, :n]
 65        c=coords3(n, x.device)
 66        bias=None
 67        if self.biased:
 68            dif=c[:,None,:]-c[None,:,:]
 69            r=torch.sqrt((dif*dif).sum(-1)+1e-10)+1e-5
 70            u=torch.sqrt((c*c).sum(-1)+1e-10)+1e-5
 71            m=torch.minimum(torch.minimum(u[:,None]/r,u[None,:]/r),torch.ones_like(r))
 72            K=r**(S-3)*m**(-SIGMA)
 73            bias=(LAMBDA*torch.log(K+1e-8)).unsqueeze(0).expand(b,-1,-1)
 74            bias=bias.repeat_interleave(HEADS, 0)
 75        weights=None
 76        for li in self.layers:
 77            z=li['norm1'](h)
 78            h2, weights=li['attn'](z,z,z, attn_mask=bias, need_weights=True, average_attn_weights=False)
 79            h=h+h2; h=h+li['ff'](li['norm2'](h))
 80        self.last_attn=weights.detach() if weights is not None else None
 81        return self.head(h.reshape(b,-1))
 82
 83
 84def seed_all(seed):
 85    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 86
 87
 88def train_one(cfg, seed, biased, capture=False):
 89    seed_all(seed)
 90    ds=get_dataset('sequence', seed, n_train=NTRAIN, n_test=NTEST)
 91    net=CoordTransformer(ds['input_shape'][0], ds['out_dim'], biased=biased)
 92    trained, metric, hist=train_model(net, ds, epochs=EPOCHS, lr=cfg['lr'], batch=128, weight_decay=cfg['weight_decay'], log=lambda *_: None)
 93    if trained is None: raise RuntimeError('bench training failed')
 94    if capture:
 95        trained.eval()
 96        with torch.no_grad():
 97            dev=next(trained.parameters()).device
 98            _=trained(ds['xte'][:64].to(dev))
 99            att=trained.last_attn.cpu().numpy()  # [B,H,N,N]
100        return float(metric), {'model': trained, 'ds': ds, 'att': att}
101    return float(metric)
102
103
104def make_train(cfg, biased):
105    return lambda seed: train_one(cfg, seed, biased)
106
107
108def main():
109    checks=math_check()
110    # Three learning rates are the shared union of both searches; baseline has no extra method knob.
111    grid=[{'lr':lr, 'weight_decay':wd} for lr,wd in
112          [(0.0015,0.0),(0.003,0.0),(0.006,0.0)]]
113    base=sweep_baseline(lambda cfg: make_train(cfg, False), grid, seeds=SWEEP_SEEDS)
114    # Same 3-point budget and same lr union for the idea; lambda is fixed a priori from stage 1.
115    idea_cfgs=grid
116    idea_sweep=[]
117    for cfg in idea_cfgs:
118        r=evaluate(make_train(cfg, True), seeds=SWEEP_SEEDS)
119        idea_sweep.append({'cfg':cfg, 'mean':r['mean']})
120    best_idea_cfg=min(idea_cfgs, key=lambda c: next(z['mean'] for z in idea_sweep if z['cfg']==c))
121    idea_full=evaluate(make_train(best_idea_cfg, True), seeds=SEEDS)
122    # Re-test trained-model attention behavior on paired seed 0, not an analytic-only signature.
123    bm, bo=train_one(base['best_cfg'], 0, False, True)
124    im, io=train_one(best_idea_cfg, 0, True, True)
125    a0=bo['att']; a1=io['att']; n=a0.shape[-1]
126    # Query endpoint (last token): compare attention odds on oldest vs latest key.
127    old, recent = float(a0[:,:,-1,0].mean()), float(a0[:,:,-1,-1].mean())
128    old_i, recent_i = float(a1[:,:,-1,0].mean()), float(a1[:,:,-1,-1].mean())
129    observed_ratio=(recent_i/(old_i+1e-12))/(recent/(old+1e-12))
130    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)
131    predicted_ratio=float(kernel_np(c[-1:], c[-1:])[0]/kernel_np(c[-1:], c[:1])[0])
132    signature={'predicted_near_vs_far_kernel_ratio':predicted_ratio,
133               'observed_attention_odds_ratio_after_bias':observed_ratio,
134               'baseline_endpoint_recent_attention':recent,
135               'idea_endpoint_recent_attention':recent_i,
136               'confirmed': bool(observed_ratio > 1.0 and observed_ratio/predicted_ratio > 1/3 and observed_ratio/predicted_ratio < 3)}
137    report=make_report('sequence','transformer_tiny',
138        {'best_cfg':base['best_cfg'],'sweep':base['sweep'],'full':base['full']}, idea_full,
139        {'mechanism_signature':signature,'idea_sweep':idea_sweep,'math_check':checks,
140         'idea_best_cfg':best_idea_cfg,'budget':{'epochs':EPOCHS,'n_train':NTRAIN,'n_test':NTEST,'seeds':list(SEEDS)}})
141    Path('bench_report.json').write_text(json.dumps(report, indent=2))
142    print(json.dumps(report, indent=2))
143
144if __name__=='__main__': main()