1import os, sys, json, random
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  6from bench import get_dataset, sweep_baseline, evaluate, make_report
  7
  8# Sequence track is the structurally matched built-in task: the representation is
  9# a multi-token feature map and the PSD feature is a token Gram/covariance.
 10DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 11EPOCHS = 12
 12NTR, NTE = 400, 200
 13BATCH = 64
 14D, M, WIN = 32, 16, 32
 15LAMBDA = 0.03
 16EPS = 1e-6
 17
 18class BottleneckNet(nn.Module):
 19    def __init__(self, kind):
 20        super().__init__()
 21        self.kind = kind
 22        self.inp = nn.Linear(1, D)
 23        self.pos = nn.Parameter(torch.zeros(1, WIN, D))
 24        nn.init.normal_(self.pos, std=.02)
 25        layer = nn.TransformerEncoderLayer(D, nhead=2, dim_feedforward=64,
 26                                            batch_first=True, dropout=0.0)
 27        self.enc = nn.TransformerEncoder(layer, 1)
 28        if kind == 'dense':
 29            self.proj = nn.Linear(WIN * D, M)
 30        else:
 31            self.a = nn.Parameter(torch.randn(M, D) / np.sqrt(D))
 32        self.head = nn.Linear(M, 1)
 33
 34    def features(self, x):
 35        h = self.inp(x.unsqueeze(-1)) + self.pos[:, :x.shape[1]]
 36        return self.enc(h)
 37
 38    def latent(self, x):
 39        h = self.features(x)
 40        if self.kind == 'dense':
 41            return self.proj(h.reshape(x.shape[0], -1))
 42        # z_i = <a_i a_i^T, HH^T/s> = mean_tokens (a_i^T h_t)^2
 43        return (torch.einsum('btd,md->btm', h, self.a).square().mean(1))
 44
 45    def forward(self, x):
 46        return self.head(self.latent(x))
 47
 48def seed_all(seed):
 49    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 50    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 51
 52def cond_loss(net, x):
 53    # Two independently perturbed views provide X,Y without changing labels.
 54    x1 = x + 0.025 * torch.randn_like(x)
 55    x2 = x + 0.025 * torch.randn_like(x)
 56    h1, h2 = net.features(x1), net.features(x2)
 57    if net.kind == 'dense':
 58        z1 = net.proj(h1.reshape(x.shape[0], -1)); z2 = net.proj(h2.reshape(x.shape[0], -1))
 59    else:
 60        z1 = torch.einsum('btd,md->btm', h1, net.a).square().mean(1)
 61        z2 = torch.einsum('btd,md->btm', h2, net.a).square().mean(1)
 62    # Frobenius norm of normalized Gram differences, vectorized over batch.
 63    X = torch.bmm(h1.transpose(1,2), h1) / WIN
 64    Y = torch.bmm(h2.transpose(1,2), h2) / WIN
 65    den = (X-Y).square().sum((1,2)).sqrt() + EPS
 66    q = (z1-z2).abs().sum(1) / den
 67    lo, hi = torch.quantile(q, .05), torch.quantile(q, .95)
 68    return torch.log(hi + EPS) - torch.log(lo + EPS)
 69
 70def train_one(seed, lr, kind, lam=0.0, return_net=False):
 71    seed_all(seed)
 72    ds = get_dataset('sequence', seed, NTR, NTE)
 73    net = BottleneckNet(kind)
 74    try:
 75        net = net.to(DEVICE); dev = DEVICE
 76        opt = torch.optim.Adam(net.parameters(), lr=lr)
 77        lossf = nn.MSELoss()
 78        xtr, ytr = ds['xtr'].to(dev), ds['ytr'].to(dev)
 79        for _ in range(EPOCHS):
 80            net.train(); p = torch.randperm(len(xtr), device=dev)
 81            for i in range(0, len(xtr), BATCH):
 82                ix = p[i:i+BATCH]
 83                pred = net(xtr[ix]); loss = lossf(pred, ytr[ix])
 84                if lam:
 85                    loss = loss + lam * cond_loss(net, xtr[ix])
 86                opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(net.parameters(), 5.0); opt.step()
 87        net.eval()
 88        with torch.no_grad():
 89            metric = float(lossf(net(ds['xte'].to(dev)), ds['yte'].to(dev)))
 90        return (metric, net, ds) if return_net else metric
 91    except RuntimeError:
 92        # Explicit CPU fallback for constrained/shared CUDA environments.
 93        seed_all(seed); net = BottleneckNet(kind).cpu(); opt = torch.optim.Adam(net.parameters(), lr=lr)
 94        xtr, ytr = ds['xtr'], ds['ytr']
 95        for _ in range(EPOCHS):
 96            p = torch.randperm(len(xtr))
 97            for i in range(0,len(xtr),BATCH):
 98                ix=p[i:i+BATCH]; loss=nn.functional.mse_loss(net(xtr[ix]),ytr[ix])
 99                if lam: loss=loss+lam*cond_loss(net,xtr[ix])
100                opt.zero_grad(); loss.backward(); opt.step()
101        with torch.no_grad(): metric=float(nn.functional.mse_loss(net(ds['xte']),ds['yte']))
102        return (metric,net,ds) if return_net else metric
103
104def stats_for_model(net, ds):
105    net.eval(); dev=next(net.parameters()).device; x=ds['xte'].to(dev)
106    with torch.no_grad():
107        x1=x+0.025*torch.randn_like(x); x2=x+0.025*torch.randn_like(x)
108        h1,h2=net.features(x1),net.features(x2)
109        if net.kind=='dense': z1=net.proj(h1.reshape(len(x),-1)); z2=net.proj(h2.reshape(len(x),-1))
110        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)
111        den=(torch.bmm(h1.transpose(1,2),h1)/WIN-torch.bmm(h2.transpose(1,2),h2)/WIN).square().sum((1,2)).sqrt()+EPS
112        q=(z1-z2).abs().sum(1)/den
113    q=q.cpu().numpy(); lo,hi=np.quantile(q,[.05,.95])
114    return {'L_q05':float(lo),'U_q95':float(hi),'beta':float(hi/max(lo,EPS)),'mean_q':float(q.mean())}
115
116def main():
117    # Search-space parity: both methods are evaluated at every lr.
118    grid=[{'lr':v} for v in (0.001,0.003,0.006)]
119    base=sweep_baseline(lambda cfg: lambda s: train_one(s,cfg['lr'],'dense'), grid)
120    best_lr=base['best_cfg']['lr']
121    idea_grid=[best_lr]+[v for v in (0.001,0.003,0.006) if v!=best_lr]
122    idea_cfg=min(idea_grid, key=lambda lr: evaluate(lambda s: train_one(s,lr,'psd',LAMBDA), seeds=(0,1,2,3))['mean'])
123    idea=evaluate(lambda s: train_one(s,idea_cfg,'psd',LAMBDA))
124    # Signature is measured from trained models, not an analytical identity.
125    bmet,bnet,bds=train_one(0,best_lr,'dense',return_net=True)
126    imet,inet,ids=train_one(0,idea_cfg,'psd',LAMBDA,return_net=True)
127    bs,ins=stats_for_model(bnet,bds),stats_for_model(inet,ids)
128    sig={'prediction':'conditioning should narrow q spread (lower beta) and raise lower distortion',
129         'baseline_seed0':bs,'idea_seed0':ins,
130         'predicted_beta_direction':ins['beta'] < bs['beta'],
131         'predicted_L_direction':ins['L_q05'] > bs['L_q05'],
132         'confirmed':bool(ins['beta'] < bs['beta'] and ins['L_q05'] > bs['L_q05'])}
133    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}})
134    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]
135    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
136    print(json.dumps(rep,indent=2))
137if __name__=='__main__': main()