1import sys, json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6import torch.nn.functional as F
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import get_dataset, train_model, evaluate, sweep_baseline, make_report
  9
 10SEED = 2951
 11EPOCHS = 8
 12NTR, NTE = 400, 200
 13LRS = [1e-3, 3e-3, 1e-2]
 14TAUS = [0.5, 1.0, 2.0]
 15
 16class DenseLayer(nn.Module):
 17    def __init__(self, d=64, tau=1.0):
 18        super().__init__(); self.tau=tau
 19        self.q=nn.Linear(d,d); self.k=nn.Linear(d,d); self.v=nn.Linear(d,d)
 20        self.o=nn.Linear(d,d); self.ff=nn.Sequential(nn.Linear(d,128),nn.ReLU(),nn.Linear(128,d))
 21        self.n1=nn.LayerNorm(d); self.n2=nn.LayerNorm(d)
 22    def forward(self,x):
 23        q,k,v=self.q(x),self.k(x),self.v(x)
 24        a=F.softmax(torch.matmul(q,k.transpose(-1,-2))/(math.sqrt(x.shape[-1])*self.tau),-1)
 25        z=self.o(torch.matmul(a,v)); return self.n2(self.n1(x+z)+self.ff(self.n1(x+z))),a
 26
 27class RoleFillerLayer(nn.Module):
 28    """Role-filler attention: positions are orthonormal roles; projected token states are fillers."""
 29    def __init__(self, win=32, d=64, tau=1.0):
 30        super().__init__(); self.win=win; self.tau=tau
 31        self.role=nn.Parameter(torch.randn(win,d)*0.02)
 32        self.qf=nn.Linear(d,d); self.ffill=nn.Linear(d,d); self.out=nn.Linear(d,d)
 33        self.ff=nn.Sequential(nn.Linear(d,128),nn.ReLU(),nn.Linear(128,d))
 34        self.n1=nn.LayerNorm(d); self.n2=nn.LayerNorm(d)
 35    def forward(self,x):
 36        # O[source, role, filler]; query at each target role specifies target role and filler.
 37        h=x + self.role[:x.shape[1]].unsqueeze(0)
 38        fillers=self.ffill(h)
 39        q=self.qf(h)
 40        scores=torch.einsum('bid,bjd->bij',q,fillers)/(math.sqrt(x.shape[-1])*self.tau)
 41        a=F.softmax(scores,-1)
 42        # Extract the target role filler from each retrieved object and rebind to target.
 43        z=torch.einsum('bij,bjd->bid',a,fillers)
 44        z=self.out(z)
 45        u=self.n1(x+z); return self.n2(u+self.ff(u)),a
 46
 47class Net(nn.Module):
 48    def __init__(self, kind='dense', win=32, d=64, tau=1.0):
 49        super().__init__(); self.kind=kind; self.win=win
 50        self.inp=nn.Linear(1,d); self.pos=nn.Parameter(torch.randn(1,win,d)*.02)
 51        Layer=DenseLayer if kind=='dense' else RoleFillerLayer
 52        self.layers=nn.ModuleList([Layer(d=d,tau=tau) if kind=='dense' else Layer(win=win,d=d,tau=tau) for _ in range(2)])
 53        self.head=nn.Linear(win*d,1); self.last_attention=None
 54    def forward(self,x):
 55        h=self.inp(x.unsqueeze(-1))+self.pos[:,:x.shape[1]]; ats=[]
 56        for layer in self.layers: h,a=layer(h); ats.append(a)
 57        self.last_attention=ats[-1].detach(); return self.head(h.reshape(x.shape[0],-1))
 58
 59def seed_all(seed):
 60    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 61    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 62
 63def run(kind, cfg, seed, return_model=False):
 64    seed_all(seed)
 65    ds=get_dataset('sequence',seed,n_train=NTR,n_test=NTE)
 66    net=Net(kind=kind,win=ds['input_shape'][0],tau=cfg['tau'])
 67    trained, metric, hist=train_model(net,ds,epochs=EPOCHS,lr=cfg['lr'],batch=128,log=lambda *_:None)
 68    if return_model: return float(metric), trained, ds
 69    return float(metric)
 70
 71def main():
 72    # Baseline sweep includes every idea-side lr and every attention temperature.
 73    grid=[{'lr':lr,'tau':tau} for lr in LRS for tau in TAUS]
 74    base=sweep_baseline(lambda c: lambda s: run('dense',c,s),grid)
 75    # Idea at the best baseline setting and two nearby settings; all are in baseline grid.
 76    idea_cfgs=[base['best_cfg'], {'lr':1e-3,'tau':1.0}, {'lr':1e-2,'tau':1.0}]
 77    idea_runs=[]
 78    for cfg in idea_cfgs:
 79        r=evaluate(lambda s,c=cfg: run('role',c,s))
 80        idea_runs.append({'cfg':cfg,'result':r})
 81    best=min(idea_runs,key=lambda z:z['result']['mean'])
 82    # NN-scale mechanism signature: measured attention entropy for learned exact-role
 83    # (diagonal target/source) versus mismatched-role (off-diagonal) queries.
 84    ent_exact=[]; ent_wrong=[]
 85    for s in range(8):
 86        metric,net,ds=run('role',best['cfg'],s,True)
 87        x=torch.as_tensor(ds['xte'][:64],dtype=torch.float32)
 88        try:
 89            device=next(net.parameters()).device
 90            with torch.no_grad(): net(x.to(device))
 91            a=net.last_attention.cpu().numpy()
 92            diag=np.arange(a.shape[1]); exact=a[:,diag,diag]
 93            wrong=a[:,diag,(diag+1)%a.shape[2]]
 94            # NN-scale signature: exact role binding should receive more mass
 95            # than a mismatched adjacent role.
 96            ent_exact.append(float(exact.mean()))
 97            ent_wrong.append(float(wrong.mean()))
 98        except Exception as e:
 99            print('signature probe failed',repr(e))
100    sig={'prediction':'role-selective exact queries should be more concentrated than mismatched queries',
101         'predicted_exact_minus_mismatched_mass':0.2,'observed_exact_role_mass':float(np.mean(ent_exact)),
102         'observed_mismatched_role_mass':float(np.mean(ent_wrong)),
103         'observed_gap':float(np.mean(ent_exact)-np.mean(ent_wrong)),
104         'confirmed':bool(ent_exact and np.mean(ent_exact)-np.mean(ent_wrong) >= 0.2),
105         'note':'probability mass is measured from attention of trained role-filler models on held-out sequence windows'}
106    rep=make_report('sequence','transformer_tiny',base,best['result'],{'idea_sweep':idea_runs,'mechanism_signature':sig})
107    # Preserve explicit mechanism_signature at top-level as required by stage-2 runner.
108    rep['mechanism_signature']=sig; rep['idea_sweep']=idea_runs
109    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
110    print(json.dumps(rep,indent=2))
111if __name__=='__main__': main()