"""Stage-2 bench for Spiderweb Hierarchical Attention. Run from this directory with the bench environment available. """ import sys, json, math, time, random from pathlib import Path import numpy as np import torch import torch.nn as nn sys.path.insert(0, "/home/maxwelhelp/all/math2nn") from bench import (get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report, count_params) SEEDS = tuple(range(8)) # This is the complete shared union: every idea lr is also a baseline lr. GRID = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 1e-2}] EPOCHS = 15 NTRAIN, NTEST = 2000, 500 class SpiderBlock(nn.Module): """Dyadic hierarchy: local leaf attention, pooled horizontal attention, then gated broadcasts from every height to descendant tokens.""" def __init__(self, d=64, heads=2, levels=5, radius=1): super().__init__() assert d % heads == 0 self.d, self.heads, self.dk = d, heads, d // heads self.levels, self.radius = levels, radius self.qkv = nn.Linear(d, 3*d) self.out = nn.Linear(d, d) self.gates = nn.Parameter(torch.full((levels + 1,), -1.0)) self.norm1 = nn.LayerNorm(d) self.norm2 = nn.LayerNorm(d) self.ff = nn.Sequential(nn.Linear(d, 128), nn.ReLU(), nn.Linear(128, d)) def _attn(self, q, k, v): # q,k,v: [B,C,D], dense only over bounded same-level neighbors. B, C, D = q.shape q = q.view(B, C, self.heads, self.dk).transpose(1, 2) k = k.view(B, C, self.heads, self.dk).transpose(1, 2) v = v.view(B, C, self.heads, self.dk).transpose(1, 2) scores = q @ k.transpose(-2, -1) / math.sqrt(self.dk) ids = torch.arange(C, device=q.device) mask = (ids[None, :] - ids[:, None]).abs() > self.radius scores = scores.masked_fill(mask[None, None], -1e9) a = torch.softmax(scores, dim=-1) return (a @ v).transpose(1, 2).reshape(B, C, D) def forward(self, x): # Residual local token attention is the fine-scale spiderweb level. z = self.norm1(x) q, k, v = self.qkv(z).chunk(3, dim=-1) x = x + self.out(self._attn(q, k, v)) B, N, D = x.shape # Mean pooling to dyadic ancestors and bounded horizontal communication. broadcast = torch.zeros_like(x) for lev in range(1, self.levels + 1): size = 2 ** lev C = (N + size - 1) // size pad = C * size - N xp = torch.cat([x, x[:, -1:, :].expand(B, pad, D)], dim=1) if pad else x cells = xp.view(B, C, size, D).mean(2) cq, ck, cv = self.qkv(cells).chunk(3, dim=-1) u = self.out(self._attn(cq, ck, cv)) expanded = u.repeat_interleave(size, dim=1)[:, :N] broadcast = broadcast + torch.sigmoid(self.gates[lev]) * expanded x = x + broadcast / max(1, self.levels) return x + self.ff(self.norm2(x)) class SpiderTransformer(nn.Module): def __init__(self, win, out_dim=1, d=64, depth=2): super().__init__() self.inp = nn.Linear(1, d) self.pos = nn.Parameter(torch.zeros(1, win, d)) nn.init.normal_(self.pos, std=.02) self.enc = nn.ModuleList([SpiderBlock(d=d, heads=2, levels=5, radius=1) for _ in range(depth)]) self.head = nn.Linear(win*d, out_dim) def forward(self, x): h = self.inp(x.unsqueeze(-1)) + self.pos[:, :x.shape[1]] for layer in self.enc: h = layer(h) return self.head(h.reshape(x.shape[0], -1)) def setup(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def train_one(kind, seed, lr): setup(seed) ds = get_dataset("sequence", seed, n_train=NTRAIN, n_test=NTEST) if kind == "baseline": net = make_model("transformer_tiny", ds["input_shape"], ds["out_dim"]) else: net = SpiderTransformer(ds["input_shape"][0], ds["out_dim"]) _, metric, _ = train_model(net, ds, epochs=EPOCHS, lr=lr, batch=128, log=lambda *_: None) return float(metric) def signature(seed, lr, kind): """Measured on trained systems: output sensitivity to first-token perturbation.""" setup(seed); ds = get_dataset("sequence", seed, n_train=NTRAIN, n_test=NTEST) if kind == "baseline": net = make_model("transformer_tiny", ds["input_shape"], 1) else: net = SpiderTransformer(ds["input_shape"][0], 1) net, _, _ = train_model(net, ds, epochs=EPOCHS, lr=lr, batch=128, log=lambda *_: None) net.eval(); device = next(net.parameters()).device; x = ds["xte"][:1].clone().to(device) with torch.no_grad(): y0 = net(x).item(); x1 = x.clone(); x1[0, 0] += 1.0; y1 = net(x1).item() x2 = x.clone(); x2[0, -1] += 1.0; y2 = net(x2).item() return {"far_first_token_effect": abs(y1-y0), "local_last_token_effect": abs(y2-y0), "far_to_local_ratio": abs(y1-y0)/(abs(y2-y0)+1e-8)} def main(): t0=time.time() # Baseline sweep on protocol seeds; all configs use exactly the idea's lr union. base = sweep_baseline(lambda cfg: (lambda s: train_one("baseline", s, cfg["lr"])), GRID) # Explicitly evaluate idea at all three shared settings on the same 8 seeds. idea_cfgs=[] for cfg in GRID: r=evaluate(lambda s, lr=cfg["lr"]: train_one("idea", s, lr), SEEDS) idea_cfgs.append((r,cfg)) idea, idea_cfg = min(idea_cfgs, key=lambda z:z[0]["mean"]) rep = make_report("sequence", "transformer_tiny", base, idea, {"prediction": "hierarchical bounded-neighbor attention preserves measurable long-range influence", "baseline_signature": signature(0, base["best_cfg"]["lr"], "baseline"), "idea_signature": signature(0, idea_cfg["lr"], "idea"), "confirmed": False, "confirmation_rule": "idea far-to-local ratio > baseline far-to-local ratio by at least 20%"}) bs=rep["mechanism_signature"]["baseline_signature"]; ins=rep["mechanism_signature"]["idea_signature"] rep["mechanism_signature"]["confirmed"] = ins["far_to_local_ratio"] >= 1.2*bs["far_to_local_ratio"] 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)), "params_idea":count_params(SpiderTransformer(32,1)), "elapsed_sec":time.time()-t0, "n_train":NTRAIN,"n_test":NTEST,"epochs":EPOCHS} Path("bench_report.json").write_text(json.dumps(rep, indent=2)) print(json.dumps(rep, indent=2)) if __name__ == "__main__": main()