1import sys, json, math, random
 2from pathlib import Path
 3import numpy as np
 4import torch
 5import torch.nn as nn
 6import torch.nn.functional as F
 7
 8sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
 9from bench import get_dataset, train_model, sweep_baseline, make_report, evaluate
10
11SEEDS = tuple(range(8)); SWEEP_SEEDS = tuple(range(4)); WIN = 32
12EPOCHS = 10; BATCH = 128; WEIGHT_DECAY = 0.0
13LRS = [0.0015, 0.003, 0.006]; M = 8; TAU_MIN = 1e-3
14CAPTURE = {}
15
16class BiasAttention(nn.Module):
17    def __init__(self, d, heads, win, kind, gamma=1.0):
18        super().__init__(); assert d % heads == 0
19        self.d, self.heads, self.dk, self.win, self.kind = d, heads, d//heads, win, kind
20        self.gamma = float(gamma); self.qkv = nn.Linear(d, 3*d); self.out = nn.Linear(d, d)
21        if kind == "table": self.bias = nn.Parameter(torch.zeros(heads, win))
22        else:
23            self.alpha = nn.Parameter(torch.zeros(heads, M))
24            init = torch.logspace(math.log10(.03), math.log10(1.0), M)
25            self.beta = nn.Parameter(torch.log(torch.expm1(init-TAU_MIN)).repeat(heads, 1))
26    def lag_bias(self, device):
27        d = torch.arange(self.win, device=device, dtype=torch.float32)
28        if self.kind == "table": return self.bias
29        w = F.softmax(self.alpha, dim=-1); tau = F.softplus(self.beta) + TAU_MIN
30        k = torch.exp(-d[None,:,None] * tau[:,None,:]).matmul(w[...,None]).squeeze(-1)
31        return torch.log(k + 1e-6) * self.gamma
32    def forward(self, x):
33        B,L,_ = x.shape; q,k,v = self.qkv(x).chunk(3, dim=-1)
34        def split(z): return z.view(B,L,self.heads,self.dk).transpose(1,2)
35        q,k,v = map(split,(q,k,v)); logits = (q @ k.transpose(-2,-1))/math.sqrt(self.dk)
36        lb = self.lag_bias(x.device)
37        lags = (torch.arange(L,device=x.device)[None,:]-torch.arange(L,device=x.device)[:,None]).clamp(min=0)
38        logits = logits + lb[:,lags][None]
39        logits = logits.masked_fill(torch.triu(torch.ones(L,L,device=x.device,dtype=torch.bool),1),-1e4)
40        a = torch.softmax(logits,dim=-1); y=(a@v).transpose(1,2).contiguous().view(B,L,self.d)
41        return self.out(y),a
42
43class Block(nn.Module):
44    def __init__(self,d,heads,win,kind,gamma):
45        super().__init__(); self.n1=nn.LayerNorm(d); self.attn=BiasAttention(d,heads,win,kind,gamma)
46        self.n2=nn.LayerNorm(d); self.ff=nn.Sequential(nn.Linear(d,128),nn.GELU(),nn.Linear(128,d))
47    def forward(self,x):
48        z,a=self.attn(self.n1(x)); x=x+z; return x+self.ff(self.n2(x)),a
49
50class LagTransformer(nn.Module):
51    def __init__(self,kind,gamma=1.0,depth=2,d=64,heads=2,win=WIN):
52        super().__init__(); self.inp=nn.Linear(1,d); self.pos=nn.Parameter(torch.zeros(1,win,d)); nn.init.normal_(self.pos,std=.02)
53        self.blocks=nn.ModuleList([Block(d,heads,win,kind,gamma) for _ in range(depth)]); self.head=nn.Linear(win*d,1); self.last_attn=None
54    def forward(self,x):
55        h=self.inp(x.unsqueeze(-1))+self.pos[:,:x.shape[1]]; aa=[]
56        for block in self.blocks: h,a=block(h); aa.append(a)
57        self.last_attn=aa[-1].detach(); return self.head(h.reshape(x.shape[0],-1))
58
59def make_train(kind,cfg,capture=False):
60    def run(seed):
61        random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
62        ds=get_dataset("sequence",seed,n_train=800,n_test=300); model=LagTransformer(kind,gamma=cfg["gamma"])
63        net,metric,_=train_model(model,ds,epochs=EPOCHS,lr=cfg["lr"],batch=BATCH,weight_decay=WEIGHT_DECAY,log=lambda *_:None)
64        if net is None: return float("inf")
65        if capture:
66            with torch.no_grad():
67                dev=next(net.parameters()).device; b=net.blocks[-1].attn.lag_bias(dev).cpu().numpy()
68                net.eval(); net(ds["xte"][:300].to(dev)); a=net.last_attn.cpu().numpy(); L=a.shape[-1]
69                av=np.zeros(L); ct=np.zeros(L)
70                for q in range(L):
71                    for k in range(q+1): av[q-k]+=a[:,:,q,k].mean(); ct[q-k]+=1
72                CAPTURE[int(seed)]={"bias":b.tolist(),"attention":(av/np.maximum(ct,1)).tolist()}
73        return float(metric)
74    return run
75
76def signature():
77    first=[]; second=[]; att=[]
78    for rec in CAPTURE.values():
79        b=np.asarray(rec["bias"]); first.extend(np.diff(b,axis=1).ravel()); second.extend(np.diff(b,n=2,axis=1).ravel()); att.append(np.diff(rec["attention"]).max())
80    f=float(np.max(first)); s=float(np.min(second))
81    return {"prediction":"trained mixture log-kernel bias decreases with lag; retained content logits can make total attention nonmonotone","observed_bias_max_first_difference":f,"observed_bias_min_second_difference":s,"observed_attention_max_first_difference":float(max(att)),"n_models":len(CAPTURE),"confirmed":bool(f<=1e-7)}
82
83def main():
84    grid=[{"lr":lr,"gamma":1.0} for lr in LRS]
85    base=sweep_baseline(lambda cfg:make_train("table",cfg),grid)
86    # Same union of learning rates; select idea on the same four sweep seeds.
87    tried=[]
88    for cfg in grid:
89        r=evaluate(make_train("mixture",cfg),seeds=SWEEP_SEEDS); tried.append({"cfg":cfg,"mean":r["mean"]})
90    best_cfg=min(grid,key=lambda c: next(x["mean"] for x in tried if x["cfg"]==c)); CAPTURE.clear()
91    idea=evaluate(make_train("mixture",best_cfg,capture=True),seeds=SEEDS)
92    report=make_report("sequence","transformer_tiny",base,idea,{"mechanism_signature":signature(),"protocol_notes":"Matched sequence track; shared 2-block d=64 causal transformer, differing only in relative lag bias: unconstrained table versus 8-positive-exponential mixture."})
93    report["idea_sweep"]=tried; report["idea_config"]=best_cfg; report["search_space"]={"baseline_grid":grid,"idea_grid":grid}
94    Path("bench_report.json").write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2))
95if __name__ == "__main__": main()