import os, sys, json, random import numpy as np import torch import torch.nn as nn sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import get_dataset, sweep_baseline, evaluate, make_report # Sequence track is the structurally matched built-in task: the representation is # a multi-token feature map and the PSD feature is a token Gram/covariance. DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' EPOCHS = 12 NTR, NTE = 400, 200 BATCH = 64 D, M, WIN = 32, 16, 32 LAMBDA = 0.03 EPS = 1e-6 class BottleneckNet(nn.Module): def __init__(self, kind): super().__init__() self.kind = kind self.inp = nn.Linear(1, D) self.pos = nn.Parameter(torch.zeros(1, WIN, D)) nn.init.normal_(self.pos, std=.02) layer = nn.TransformerEncoderLayer(D, nhead=2, dim_feedforward=64, batch_first=True, dropout=0.0) self.enc = nn.TransformerEncoder(layer, 1) if kind == 'dense': self.proj = nn.Linear(WIN * D, M) else: self.a = nn.Parameter(torch.randn(M, D) / np.sqrt(D)) self.head = nn.Linear(M, 1) def features(self, x): h = self.inp(x.unsqueeze(-1)) + self.pos[:, :x.shape[1]] return self.enc(h) def latent(self, x): h = self.features(x) if self.kind == 'dense': return self.proj(h.reshape(x.shape[0], -1)) # z_i = = mean_tokens (a_i^T h_t)^2 return (torch.einsum('btd,md->btm', h, self.a).square().mean(1)) def forward(self, x): return self.head(self.latent(x)) def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def cond_loss(net, x): # Two independently perturbed views provide X,Y without changing labels. x1 = x + 0.025 * torch.randn_like(x) x2 = x + 0.025 * torch.randn_like(x) h1, h2 = net.features(x1), net.features(x2) if net.kind == 'dense': z1 = net.proj(h1.reshape(x.shape[0], -1)); z2 = net.proj(h2.reshape(x.shape[0], -1)) else: z1 = torch.einsum('btd,md->btm', h1, net.a).square().mean(1) z2 = torch.einsum('btd,md->btm', h2, net.a).square().mean(1) # Frobenius norm of normalized Gram differences, vectorized over batch. X = torch.bmm(h1.transpose(1,2), h1) / WIN Y = torch.bmm(h2.transpose(1,2), h2) / WIN den = (X-Y).square().sum((1,2)).sqrt() + EPS q = (z1-z2).abs().sum(1) / den lo, hi = torch.quantile(q, .05), torch.quantile(q, .95) return torch.log(hi + EPS) - torch.log(lo + EPS) def train_one(seed, lr, kind, lam=0.0, return_net=False): seed_all(seed) ds = get_dataset('sequence', seed, NTR, NTE) net = BottleneckNet(kind) try: net = net.to(DEVICE); dev = DEVICE opt = torch.optim.Adam(net.parameters(), lr=lr) lossf = nn.MSELoss() xtr, ytr = ds['xtr'].to(dev), ds['ytr'].to(dev) for _ in range(EPOCHS): net.train(); p = torch.randperm(len(xtr), device=dev) for i in range(0, len(xtr), BATCH): ix = p[i:i+BATCH] pred = net(xtr[ix]); loss = lossf(pred, ytr[ix]) if lam: loss = loss + lam * cond_loss(net, xtr[ix]) opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(net.parameters(), 5.0); opt.step() net.eval() with torch.no_grad(): metric = float(lossf(net(ds['xte'].to(dev)), ds['yte'].to(dev))) return (metric, net, ds) if return_net else metric except RuntimeError: # Explicit CPU fallback for constrained/shared CUDA environments. seed_all(seed); net = BottleneckNet(kind).cpu(); opt = torch.optim.Adam(net.parameters(), lr=lr) xtr, ytr = ds['xtr'], ds['ytr'] for _ in range(EPOCHS): p = torch.randperm(len(xtr)) for i in range(0,len(xtr),BATCH): ix=p[i:i+BATCH]; loss=nn.functional.mse_loss(net(xtr[ix]),ytr[ix]) if lam: loss=loss+lam*cond_loss(net,xtr[ix]) opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): metric=float(nn.functional.mse_loss(net(ds['xte']),ds['yte'])) return (metric,net,ds) if return_net else metric def stats_for_model(net, ds): net.eval(); dev=next(net.parameters()).device; x=ds['xte'].to(dev) with torch.no_grad(): x1=x+0.025*torch.randn_like(x); x2=x+0.025*torch.randn_like(x) h1,h2=net.features(x1),net.features(x2) if net.kind=='dense': z1=net.proj(h1.reshape(len(x),-1)); z2=net.proj(h2.reshape(len(x),-1)) else: z1=torch.einsum('btd,md->btm',h1,net.a).square().mean(1); z2=torch.einsum('btd,md->btm',h2,net.a).square().mean(1) den=(torch.bmm(h1.transpose(1,2),h1)/WIN-torch.bmm(h2.transpose(1,2),h2)/WIN).square().sum((1,2)).sqrt()+EPS q=(z1-z2).abs().sum(1)/den q=q.cpu().numpy(); lo,hi=np.quantile(q,[.05,.95]) return {'L_q05':float(lo),'U_q95':float(hi),'beta':float(hi/max(lo,EPS)),'mean_q':float(q.mean())} def main(): # Search-space parity: both methods are evaluated at every lr. grid=[{'lr':v} for v in (0.001,0.003,0.006)] base=sweep_baseline(lambda cfg: lambda s: train_one(s,cfg['lr'],'dense'), grid) best_lr=base['best_cfg']['lr'] idea_grid=[best_lr]+[v for v in (0.001,0.003,0.006) if v!=best_lr] idea_cfg=min(idea_grid, key=lambda lr: evaluate(lambda s: train_one(s,lr,'psd',LAMBDA), seeds=(0,1,2,3))['mean']) idea=evaluate(lambda s: train_one(s,idea_cfg,'psd',LAMBDA)) # Signature is measured from trained models, not an analytical identity. bmet,bnet,bds=train_one(0,best_lr,'dense',return_net=True) imet,inet,ids=train_one(0,idea_cfg,'psd',LAMBDA,return_net=True) bs,ins=stats_for_model(bnet,bds),stats_for_model(inet,ids) sig={'prediction':'conditioning should narrow q spread (lower beta) and raise lower distortion', 'baseline_seed0':bs,'idea_seed0':ins, 'predicted_beta_direction':ins['beta'] < bs['beta'], 'predicted_L_direction':ins['L_q05'] > bs['L_q05'], 'confirmed':bool(ins['beta'] < bs['beta'] and ins['L_q05'] > bs['L_q05'])} rep=make_report('sequence','transformer_tiny',base,idea,{'track_match':'multi-token sequence correlations','signature':sig,'hyperparameters':{'epochs':EPOCHS,'m':M,'d':D,'lambda':LAMBDA,'idea_lr':idea_cfg}}) rep['idea_sweep']=[{'lr':lr,'mean_4seed':evaluate(lambda s,lr=lr:train_one(s,lr,'psd',LAMBDA),seeds=(0,1,2,3))['mean']} for lr in idea_grid] with open('bench_report.json','w') as f: json.dump(rep,f,indent=2) print(json.dumps(rep,indent=2)) if __name__=='__main__': main()