1import os, json, random
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5import sys
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, train_model, make_report
  8from bench.protocol import evaluate, sweep_baseline
  9from bench.models import transformer_tiny
 10
 11TRACK='sequence'; MODEL='transformer_tiny'; SEEDS=tuple(range(8))
 12# Small, equal-budget benchmark: all candidate lrs are evaluated by both methods.
 13NTR,NTE,EPOCHS,BATCH=400,400,5,128
 14LR_GRID=[1e-3,3e-3,1e-2]
 15LAMBDA=0.12; RHO=0.5; M=32
 16
 17def seed_all(s):
 18    random.seed(s); np.random.seed(s); torch.manual_seed(s)
 19    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
 20
 21def make_net(d):
 22    # Same base transformer as bench, with the sigmoid propensity encoder required
 23    # by the idea; both systems use this identical architecture.
 24    return SigmoidTransformer(d['input_shape'][0], d['out_dim'])
 25
 26class SigmoidTransformer(nn.Module):
 27    def __init__(self, win, out_dim, dim=64):
 28        super().__init__()
 29        self.inp=nn.Linear(1,dim)
 30        self.pos=nn.Parameter(torch.zeros(1,win,dim)); nn.init.normal_(self.pos,std=.02)
 31        layer=nn.TransformerEncoderLayer(dim,nhead=2,dim_feedforward=128,batch_first=True,dropout=0.)
 32        self.enc=nn.TransformerEncoder(layer,2)
 33        self.head=nn.Linear(win*dim,out_dim)
 34    def forward(self,x):
 35        h=torch.sigmoid(self.inp(x.unsqueeze(-1))) + self.pos[:,:x.shape[1]]
 36        return self.head(self.enc(h).reshape(x.shape[0],-1))
 37    def propensity(self,x):
 38        return torch.sigmoid(self.inp(x.unsqueeze(-1))).mean(1)
 39
 40def baseline_run(cfg, seed, keep=False):
 41    seed_all(seed); d=get_dataset(TRACK,seed,n_train=NTR,n_test=NTE)
 42    net=make_net(d)
 43    # train_model is the canonical baseline path; model architecture is shared.
 44    _, metric, _=train_model(net,d,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH,log=lambda *_:None)
 45    return float(metric)
 46
 47def local_run(cfg, seed, collect=False, device_override=None):
 48    seed_all(seed); d=get_dataset(TRACK,seed,n_train=NTR,n_test=NTE)
 49    net=make_net(d)
 50    # Adam trains the non-encoder parameters identically; the intervention is
 51    # exclusively the online fixed-projection update of inp.weight/bias.
 52    device=device_override or ('cuda' if torch.cuda.is_available() else 'cpu')
 53    try:
 54        net=net.to(device); xtr,ytr=d['xtr'].to(device),d['ytr'].to(device)
 55        xte,yte=d['xte'].to(device),d['yte'].to(device)
 56        opt=torch.optim.Adam([p for n,p in net.named_parameters() if not n.startswith('inp.')],lr=cfg['lr'])
 57        rng=torch.Generator(device=device); rng.manual_seed(seed+991)
 58        A=torch.randn(M,64,device=device,generator=rng)/np.sqrt(M)
 59        before=net.inp.weight.detach().clone(); sign_num=sign_den=0.
 60        lossf=nn.M1Loss() if False else nn.MSELoss()
 61        for ep in range(EPOCHS):
 62            net.train(); perm=torch.randperm(len(xtr),device=device)
 63            for st in range(0,len(xtr),BATCH):
 64                ix=perm[st:st+BATCH]; pred=net(xtr[ix]); loss=lossf(pred,ytr[ix])
 65                opt.zero_grad(); loss.backward(); opt.step()
 66                # Ordered consecutive examples; vectorized local rank-one updates.
 67                if st == 0 or len(ix) > 1:
 68                    xa, xb = xtr[ix[:-1]], xtr[ix[1:]]
 69                    with torch.no_grad():
 70                        ha = torch.sigmoid(net.inp(xa.unsqueeze(-1))).mean(1)
 71                        hb = torch.sigmoid(net.inp(xb.unsqueeze(-1))).mean(1)
 72                        drive = (hb-ha) @ A.t() @ A
 73                        u = drive - LAMBDA*(ha-RHO)
 74                        delta = cfg['lr'] * ha*(1-ha)*u
 75                        net.inp.weight.add_(delta.t() @ xa.mean(1, keepdim=True) / max(1,len(ix)-1))
 76                        net.inp.bias.add_(delta.mean(0))
 77                        sign_num += float((delta*drive).sum())
 78                        sign_den += float(delta.abs().sum())
 79        net.eval()
 80        with torch.no_grad(): metric=float(((net(xte)-yte)**2).mean())
 81        if collect:
 82            with torch.no_grad():
 83                h=net.propensity(xte).flatten().cpu().numpy()
 84            return metric, {'mean_hidden':float(h.mean()),'saturated_fraction':float(np.mean((h<.05)|(h>.95))),
 85                            'projection_alignment':float(sign_num/(sign_den+1e-12)),'device':device}
 86        return metric
 87    except RuntimeError:
 88        # Explicit one-way GPU fallback; do not recurse if CPU itself fails.
 89        if device == 'cuda':
 90            torch.cuda.empty_cache()
 91            return local_run(cfg, seed, collect, device_override='cpu')
 92        raise
 93
 94def main():
 95    grid=[{'lr':x} for x in LR_GRID]
 96    base=sweep_baseline(lambda c:(lambda s: baseline_run(c,s)),grid)
 97    best=base['best_cfg']
 98    idea_grid=grid
 99    idea_sweep=[{'cfg':c,'mean':evaluate(lambda s,c=c:local_run(c,s),SEEDS)['mean']} for c in idea_grid]
100    best_idea=min(idea_sweep,key=lambda z:z['mean'])['cfg']
101    idea=evaluate(lambda s:local_run(best_idea,s),SEEDS)
102    sig=[]
103    for s in SEEDS:
104        _,st=local_run(best_idea,s,True); sig.append(st)
105    sigmean={k:float(np.mean([z[k] for z in sig])) for k in sig[0] if k!='device'}
106    sigmean['per_seed']=sig; sigmean['predicted_projection_alignment_positive']=True
107    sigmean['confirmed']=bool(sigmean['projection_alignment']>0)
108    base['idea_grid']=idea_sweep
109    rep=make_report(TRACK,MODEL,base,idea,extra=sigmean)
110    rep['protocol_notes']={'architecture_match':'identical sigmoid-input transformer; only inp training rule differs',
111      'track_justification':'sequence forecast contains ordered multi-token temporal windows',
112      'budget':{'n_train':NTR,'n_test':NTE,'epochs':EPOCHS,'batch':BATCH,'lr_union':LR_GRID},
113      'custom_track':None}
114    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
115    print(json.dumps(rep,indent=2))
116if __name__=='__main__': main()