1import json, math, random, sys
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import train_model, evaluate, sweep_baseline, make_report, get_dataset
  8
  9SEEDS = tuple(range(8))
 10NTR, NTE, EPOCHS = 400, 200, 20
 11LRS = [1e-3, 3e-3, 1e-2]
 12WHITEN_DELTAS = [1e-4]
 13
 14
 15def seed_all(seed):
 16    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 17    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 18
 19
 20def prep(seed):
 21    d = get_dataset('sequence', seed=seed, n_train=NTR, n_test=NTE)
 22    d = dict(d)
 23    for k in ('xtr','xte'):
 24        d[k] = torch.as_tensor(d[k], dtype=torch.float32)
 25    for k in ('ytr','yte'):
 26        d[k] = torch.as_tensor(d[k], dtype=torch.float32)
 27    d['ytr'] = d['ytr'].reshape(-1); d['yte'] = d['yte'].reshape(-1)
 28    d['task'] = 'regression'; d['metric'] = 'mse'
 29    d['input_shape'] = tuple(d['xtr'].shape[1:])
 30    return d
 31
 32
 33class OTAttention(nn.Module):
 34    """Small differentiable entropic OT attention over all window tokens.
 35
 36    baseline learns raw pair-feature coefficients. idea applies the fixed
 37    row/column quotient projector followed by empirical covariance whitening;
 38    both variants share dimensions, encoder, Sinkhorn iterations, and readout.
 39    """
 40    def __init__(self, d, idea=False, delta=1e-4, eps=.15, iters=12):
 41        super().__init__(); self.d=d; self.idea=idea; self.eps=eps; self.iters=iters
 42        self.q = nn.Linear(d, d, bias=False); self.k = nn.Linear(d, d, bias=False)
 43        self.v = nn.Linear(d, d, bias=False)
 44        self.theta = nn.Parameter(torch.randn(3) * .03)
 45        self.delta = delta
 46        # State-independent feature covariance estimate is updated from each
 47        # input batch and used only in the forward parameterization.
 48        self.register_buffer('cov', torch.eye(3))
 49        self.register_buffer('seen', torch.tensor(False))
 50
 51    def pair_features(self, z):
 52        q, k = self.q(z), self.k(z)
 53        # Features are scalar pair interaction plus row and column nuisance
 54        # coordinates, making the gauge structure explicit.
 55        score = torch.einsum('bid,bjd->bij', q, k) / math.sqrt(self.d)
 56        row = q.mean(-1, keepdim=True).expand(-1, -1, z.shape[1])
 57        col = k.mean(-1, keepdim=True).transpose(1,2).expand(-1, z.shape[1], -1)
 58        return torch.stack((score, row, col), -1), q, k
 59
 60    def forward(self, z, return_plan=False):
 61        phi, q, k = self.pair_features(z)
 62        b,n,_,f = phi.shape
 63        flat = phi.reshape(-1, f)
 64        with torch.no_grad():
 65            c = (flat.T @ flat) / max(flat.shape[0], 1)
 66            if not bool(self.seen): self.cov.copy_(c + 1e-6*torch.eye(f, device=c.device)); self.seen.fill_(True)
 67            else: self.cov.mul_(.95).add_(.05*c)
 68        if self.idea:
 69            # Double-centering is the orthogonal complement of row+column
 70            # potentials; covariance whitening is in the retained coordinates.
 71            a = phi - phi.mean(2, keepdim=True) - phi.mean(1, keepdim=True) + phi.mean((1,2), keepdim=True)
 72            c = self.cov + self.delta*torch.eye(f, device=z.device)
 73            e,u = torch.linalg.eigh(c)
 74            w = (u * torch.rsqrt(e.clamp_min(1e-8))) @ u.T
 75            a = torch.einsum('bijk,kl->bijl', a, w)
 76        else: a = phi
 77        cost = -torch.einsum('f,bijn->bij', self.theta, a)
 78        logk = -cost / self.eps
 79        lu = torch.zeros((b,n), device=z.device); lv = torch.zeros((b,n), device=z.device)
 80        norm = -math.log(n)
 81        for _ in range(self.iters):
 82            lu = norm - torch.logsumexp(logk + lv[:,None,:], 2)
 83            lv = norm - torch.logsumexp(logk + lu[:,:,None], 1)
 84        plan = torch.exp(lu[:,:,None] + logk + lv[:,None,:])
 85        out = torch.einsum('bij,bjd->bid', plan*n, self.v(z))
 86        pooled = out.mean(1) + z.mean(1)
 87        return (pooled, plan) if return_plan else pooled
 88
 89
 90class TinySequence(nn.Module):
 91    def __init__(self, input_shape, idea=False, delta=1e-4):
 92        super().__init__(); d=int(input_shape[-1])
 93        self.inp=nn.Linear(1,32); self.ot=OTAttention(32,idea,delta)
 94        self.head=nn.Sequential(nn.LayerNorm(32),nn.Linear(32,1))
 95    def forward(self,x): return self.head(self.ot(torch.tanh(self.inp(x.unsqueeze(-1))))).squeeze(-1)
 96
 97
 98def run(kind, lr, seed, delta=1e-4):
 99    seed_all(seed); ds=prep(seed)
100    model=TinySequence(ds['input_shape'], kind=='idea', delta)
101    net, metric, hist=train_model(model, ds, epochs=EPOCHS, lr=lr, batch=128, log=lambda *a,**k:None)
102    return float(metric)
103
104
105def base_factory(cfg): return lambda seed: run('baseline', float(cfg['lr']), seed)
106def idea_factory(cfg): return lambda seed: run('idea', float(cfg['lr']), seed, float(cfg['delta']))
107
108
109def mechanism_signature():
110    seed_all(9017); ds=prep(9017); d=int(ds['input_shape'][-1]);
111    # Trained-model behavior: compare plan changes under row/column nuisance
112    # perturbations and an identifiable interaction perturbation.
113    vals=[]
114    for kind in ('baseline','idea'):
115        seed_all(9017); m=TinySequence(ds['input_shape'],kind=='idea',1e-4)
116        m,_,_=train_model(m,ds,epochs=EPOCHS,lr=3e-3,batch=128,log=lambda *a,**k:None)
117        dev=next(m.parameters()).device; x=ds['xte'][:16].to(dev); z=torch.tanh(m.inp(x.unsqueeze(-1)))
118        with torch.no_grad():
119            _,p=m.ot(z,True); zrow=z.clone(); zrow[:,:,0]+=0.1
120            _,pr=m.ot(zrow,True); zint=z.clone(); zint[:,0,0]+=0.1
121            _,pi=m.ot(zint,True)
122        vals.append({'kind':kind,'rowcol_plan_rms':float((p-pr).pow(2).mean().sqrt()),'interaction_plan_rms':float((p-pi).pow(2).mean().sqrt())})
123    ratio=vals[1]['rowcol_plan_rms']/(vals[0]['rowcol_plan_rms']+1e-12)
124    return {'prediction':'quotient attention suppresses row/column nuisance sensitivity while retaining interaction sensitivity','baseline':vals[0],'idea':vals[1],'predicted_ratio_bound':1.0,'observed_ratio':ratio,'confirmed':bool(np.isfinite(ratio) and ratio<=1.05)}
125
126
127def main():
128    grid=[{'lr':lr,'delta':delta} for lr in LRS for delta in WHITEN_DELTAS]
129    # Search-space parity: baseline evaluates every lr/delta pair too, though
130    # delta is inert there, so the union of method settings is identical.
131    base=sweep_baseline(base_factory,grid,seeds=SEEDS)
132    trials=[{'cfg':c,'result':evaluate(idea_factory(c),SEEDS)} for c in grid]
133    best=min(trials,key=lambda q:q['result']['mean'])
134    rep=make_report('sequence','transformer_tiny',base,best['result'],{'idea_config':best['cfg'],'idea_sweep':trials,'mechanism_signature':mechanism_signature()})
135    rep['mechanism_signature']=rep.pop('mechanism_signature')
136    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
137    print(json.dumps(rep,indent=2))
138
139if __name__=='__main__': main()