1"""Stage-2 bench for Spiderweb Hierarchical Attention.
  2Run from this directory with the bench environment available.
  3"""
  4import sys, json, math, time, random
  5from pathlib import Path
  6import numpy as np
  7import torch
  8import torch.nn as nn
  9
 10sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
 11from bench import (get_dataset, make_model, train_model, evaluate,
 12                   sweep_baseline, make_report, count_params)
 13
 14SEEDS = tuple(range(8))
 15# This is the complete shared union: every idea lr is also a baseline lr.
 16GRID = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 1e-2}]
 17EPOCHS = 15
 18NTRAIN, NTEST = 2000, 500
 19
 20class SpiderBlock(nn.Module):
 21    """Dyadic hierarchy: local leaf attention, pooled horizontal attention,
 22    then gated broadcasts from every height to descendant tokens."""
 23    def __init__(self, d=64, heads=2, levels=5, radius=1):
 24        super().__init__()
 25        assert d % heads == 0
 26        self.d, self.heads, self.dk = d, heads, d // heads
 27        self.levels, self.radius = levels, radius
 28        self.qkv = nn.Linear(d, 3*d)
 29        self.out = nn.Linear(d, d)
 30        self.gates = nn.Parameter(torch.full((levels + 1,), -1.0))
 31        self.norm1 = nn.LayerNorm(d)
 32        self.norm2 = nn.LayerNorm(d)
 33        self.ff = nn.Sequential(nn.Linear(d, 128), nn.ReLU(), nn.Linear(128, d))
 34
 35    def _attn(self, q, k, v):
 36        # q,k,v: [B,C,D], dense only over bounded same-level neighbors.
 37        B, C, D = q.shape
 38        q = q.view(B, C, self.heads, self.dk).transpose(1, 2)
 39        k = k.view(B, C, self.heads, self.dk).transpose(1, 2)
 40        v = v.view(B, C, self.heads, self.dk).transpose(1, 2)
 41        scores = q @ k.transpose(-2, -1) / math.sqrt(self.dk)
 42        ids = torch.arange(C, device=q.device)
 43        mask = (ids[None, :] - ids[:, None]).abs() > self.radius
 44        scores = scores.masked_fill(mask[None, None], -1e9)
 45        a = torch.softmax(scores, dim=-1)
 46        return (a @ v).transpose(1, 2).reshape(B, C, D)
 47
 48    def forward(self, x):
 49        # Residual local token attention is the fine-scale spiderweb level.
 50        z = self.norm1(x)
 51        q, k, v = self.qkv(z).chunk(3, dim=-1)
 52        x = x + self.out(self._attn(q, k, v))
 53        B, N, D = x.shape
 54        # Mean pooling to dyadic ancestors and bounded horizontal communication.
 55        broadcast = torch.zeros_like(x)
 56        for lev in range(1, self.levels + 1):
 57            size = 2 ** lev
 58            C = (N + size - 1) // size
 59            pad = C * size - N
 60            xp = torch.cat([x, x[:, -1:, :].expand(B, pad, D)], dim=1) if pad else x
 61            cells = xp.view(B, C, size, D).mean(2)
 62            cq, ck, cv = self.qkv(cells).chunk(3, dim=-1)
 63            u = self.out(self._attn(cq, ck, cv))
 64            expanded = u.repeat_interleave(size, dim=1)[:, :N]
 65            broadcast = broadcast + torch.sigmoid(self.gates[lev]) * expanded
 66        x = x + broadcast / max(1, self.levels)
 67        return x + self.ff(self.norm2(x))
 68
 69class SpiderTransformer(nn.Module):
 70    def __init__(self, win, out_dim=1, d=64, depth=2):
 71        super().__init__()
 72        self.inp = nn.Linear(1, d)
 73        self.pos = nn.Parameter(torch.zeros(1, win, d))
 74        nn.init.normal_(self.pos, std=.02)
 75        self.enc = nn.ModuleList([SpiderBlock(d=d, heads=2, levels=5, radius=1) for _ in range(depth)])
 76        self.head = nn.Linear(win*d, out_dim)
 77    def forward(self, x):
 78        h = self.inp(x.unsqueeze(-1)) + self.pos[:, :x.shape[1]]
 79        for layer in self.enc:
 80            h = layer(h)
 81        return self.head(h.reshape(x.shape[0], -1))
 82
 83def setup(seed):
 84    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 85    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 86
 87def train_one(kind, seed, lr):
 88    setup(seed)
 89    ds = get_dataset("sequence", seed, n_train=NTRAIN, n_test=NTEST)
 90    if kind == "baseline":
 91        net = make_model("transformer_tiny", ds["input_shape"], ds["out_dim"])
 92    else:
 93        net = SpiderTransformer(ds["input_shape"][0], ds["out_dim"])
 94    _, metric, _ = train_model(net, ds, epochs=EPOCHS, lr=lr, batch=128, log=lambda *_: None)
 95    return float(metric)
 96
 97def signature(seed, lr, kind):
 98    """Measured on trained systems: output sensitivity to first-token perturbation."""
 99    setup(seed); ds = get_dataset("sequence", seed, n_train=NTRAIN, n_test=NTEST)
100    if kind == "baseline": net = make_model("transformer_tiny", ds["input_shape"], 1)
101    else: net = SpiderTransformer(ds["input_shape"][0], 1)
102    net, _, _ = train_model(net, ds, epochs=EPOCHS, lr=lr, batch=128, log=lambda *_: None)
103    net.eval(); device = next(net.parameters()).device; x = ds["xte"][:1].clone().to(device)
104    with torch.no_grad():
105        y0 = net(x).item(); x1 = x.clone(); x1[0, 0] += 1.0; y1 = net(x1).item()
106        x2 = x.clone(); x2[0, -1] += 1.0; y2 = net(x2).item()
107    return {"far_first_token_effect": abs(y1-y0), "local_last_token_effect": abs(y2-y0),
108            "far_to_local_ratio": abs(y1-y0)/(abs(y2-y0)+1e-8)}
109
110def main():
111    t0=time.time()
112    # Baseline sweep on protocol seeds; all configs use exactly the idea's lr union.
113    base = sweep_baseline(lambda cfg: (lambda s: train_one("baseline", s, cfg["lr"])), GRID)
114    # Explicitly evaluate idea at all three shared settings on the same 8 seeds.
115    idea_cfgs=[]
116    for cfg in GRID:
117        r=evaluate(lambda s, lr=cfg["lr"]: train_one("idea", s, lr), SEEDS)
118        idea_cfgs.append((r,cfg))
119    idea, idea_cfg = min(idea_cfgs, key=lambda z:z[0]["mean"])
120    rep = make_report("sequence", "transformer_tiny", base, idea,
121        {"prediction": "hierarchical bounded-neighbor attention preserves measurable long-range influence",
122         "baseline_signature": signature(0, base["best_cfg"]["lr"], "baseline"),
123         "idea_signature": signature(0, idea_cfg["lr"], "idea"),
124         "confirmed": False,
125         "confirmation_rule": "idea far-to-local ratio > baseline far-to-local ratio by at least 20%"})
126    bs=rep["mechanism_signature"]["baseline_signature"]; ins=rep["mechanism_signature"]["idea_signature"]
127    rep["mechanism_signature"]["confirmed"] = ins["far_to_local_ratio"] >= 1.2*bs["far_to_local_ratio"]
128    rep["extra"]={"idea_sweep":[{"cfg":c,"result":r} for r,c in idea_cfgs],"idea_best_cfg":idea_cfg, "params_baseline":count_params(make_model("transformer_tiny", (32,), 1)),
129                   "params_idea":count_params(SpiderTransformer(32,1)), "elapsed_sec":time.time()-t0,
130                   "n_train":NTRAIN,"n_test":NTEST,"epochs":EPOCHS}
131    Path("bench_report.json").write_text(json.dumps(rep, indent=2))
132    print(json.dumps(rep, indent=2))
133
134if __name__ == "__main__": main()