1import sys, json, 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, make_report, sweep_baseline, evaluate
  9
 10SEED = 2486
 11EPOCHS = 8
 12NTR, NTE = 1200, 400
 13LRS = [1e-3, 3e-3, 6e-3]
 14A = 8
 15D = 64
 16
 17def seed_all(s):
 18    random.seed(s); np.random.seed(s); torch.manual_seed(s)
 19    if torch.cuda.is_available():
 20        try: torch.cuda.manual_seed_all(s)
 21        except Exception: pass
 22
 23def gale_check():
 24    rng = np.random.default_rng(SEED)
 25    rows=[]
 26    for a,b in [(3,7),(8,32),(12,32)]:
 27        X=rng.normal(size=(a,b))
 28        U,S,Vh=np.linalg.svd(X.T, full_matrices=True)
 29        Y=U[:,a:].T
 30        rows.append(float(np.linalg.norm(Y@X.T)))
 31    return {'sizes': [[3,7],[8,32],[12,32]], 'residuals': rows,
 32            'max_residual': max(rows), 'passed': bool(max(rows)<1e-10)}
 33
 34class BaseTransformer(nn.Module):
 35    def __init__(self, win=32, d=64, depth=2):
 36        super().__init__(); self.win=win; self.d=d
 37        self.inp=nn.Linear(1,d)
 38        self.pos=nn.Parameter(torch.zeros(1,win,d)); nn.init.normal_(self.pos,std=.02)
 39        layer=nn.TransformerEncoderLayer(d,nhead=2,dim_feedforward=128,batch_first=True,dropout=0.)
 40        self.enc=nn.TransformerEncoder(layer,depth)
 41        self.head=nn.Linear(win*d,1)
 42    def encode(self,x):
 43        h=self.inp(x.unsqueeze(-1))+self.pos[:,:x.shape[1]]
 44        return self.enc(h)
 45    def forward(self,x): return self.head(self.encode(x).reshape(x.shape[0],-1))
 46
 47class GaleTransformer(BaseTransformer):
 48    def __init__(self, win=32, d=64, a=A, depth=2):
 49        super().__init__(win,d,depth); self.a=a; self.k=win-a
 50        self.wx=nn.Linear(d,a,bias=False)
 51        self.wy=nn.Linear(d,d,bias=False)
 52        self.dual_out=nn.Linear(d,d,bias=False)
 53        self.gate=nn.Parameter(torch.tensor(0.0))
 54    def forward_features(self,x, return_stats=False):
 55        h=self.encode(x)
 56        X=self.wx(h).transpose(1,2) # B,a,T; X^T is B,T,a
 57        # full U of X^T has T columns; last T-a rows of Y span null(X)
 58        U,S,Vh=torch.linalg.svd(X.transpose(1,2), full_matrices=True)
 59        Y=U[:,:,self.a:].transpose(1,2) # B,k,T
 60        dual=Y @ self.wy(h) # B,k,d
 61        correction=Y.transpose(1,2) @ self.dual_out(dual) # B,T,d
 62        z=h + torch.sigmoid(self.gate)*correction
 63        if return_stats:
 64            # Behavior signature: overlap of learned primary and dual summaries.
 65            p=X.mean(dim=(1,2)); q=dual.mean(dim=(1,2))
 66            pc=p-p.mean(); qc=q-q.mean()
 67            corr=(pc*qc).sum()/(torch.sqrt((pc.square().sum()+1e-12)*(qc.square().sum()+1e-12)))
 68            residual=(Y @ X.transpose(1,2)).norm(dim=(1,2)).mean()
 69            return z, {'abs_summary_corr': float(corr.abs().detach()),
 70                       'null_residual': float(residual.detach()),
 71                       'gate': float(torch.sigmoid(self.gate).detach())}
 72        return z
 73    def forward(self,x): return self.head(self.forward_features(x).reshape(x.shape[0],-1))
 74
 75def dataset(seed): return get_dataset('sequence', seed, n_train=NTR, n_test=NTE)
 76
 77def run_one(kind, cfg, seed, want_stats=False):
 78    seed_all(seed); ds=dataset(seed)
 79    net=BaseTransformer() if kind=='baseline' else GaleTransformer()
 80    net, metric, hist=train_model(net, ds, epochs=EPOCHS, lr=float(cfg['lr']), batch=128)
 81    if net is None: raise RuntimeError('training failed')
 82    st=None
 83    if want_stats:
 84        net.eval()
 85        with torch.no_grad():
 86            try: _,st=net.forward_features(ds['xte'], True)
 87            except RuntimeError:
 88                # move to CPU if a shared CUDA slot failed during signature collection
 89                net=net.cpu(); _,st=net.forward_features(ds['xte'].cpu(), True)
 90    return float(metric), st
 91
 92def main():
 93    check=gale_check()
 94    assert check['passed'], check
 95    grid=[{'lr':v} for v in LRS]
 96    def mkbase(cfg): return lambda s: run_one('baseline',cfg,s)[0]
 97    base=sweep_baseline(mkbase,grid)
 98    # Same union of learning rates on idea side; best is selected on the same four sweep seeds.
 99    idea_sweep=[]
100    for cfg in grid:
101        r=evaluate(lambda s: run_one('idea',cfg,s)[0], seeds=(0,1,2,3))
102        idea_sweep.append({'cfg':cfg,'mean':r['mean']})
103    best_cfg=min(idea_sweep,key=lambda r:r['mean'])['cfg']
104    idea=evaluate(lambda s: run_one('idea',best_cfg,s)[0])
105    stats=[]
106    for s in range(8): stats.append(run_one('idea',best_cfg,s,True)[1])
107    sig={'prediction': 'dual summary should have low linear overlap with primary summary (abs corr < 0.2), while Gale residual is numerical roundoff',
108         'predicted_abs_summary_corr_max':0.2,
109         'observed_abs_summary_corr_mean':float(np.mean([x['abs_summary_corr'] for x in stats])),
110         'observed_abs_summary_corr_per_seed':[x['abs_summary_corr'] for x in stats],
111         'observed_null_residual_mean':float(np.mean([x['null_residual'] for x in stats])),
112         'observed_gate_mean':float(np.mean([x['gate'] for x in stats])),
113         'confirmed':bool(float(np.mean([x['abs_summary_corr'] for x in stats]))<0.2)}
114    report=make_report('sequence','transformer_tiny',
115        {'best_cfg':base['best_cfg'],'sweep':base['sweep'],'full':base['full']},idea,sig)
116    report['idea_sweep']=idea_sweep
117    report['core_math_check']=check
118    report['setup']={'epochs':EPOCHS,'n_train':NTR,'n_test':NTE,'a':A,'window':32,
119                     'track_justification':'sequence is structurally matched because Gale mixing operates across multi-token windows'}
120    Path('bench_report.json').write_text(json.dumps(report,indent=2))
121    print(json.dumps(report,indent=2))
122
123if __name__=='__main__': main()