Gated Local-Global Graph Attention / stage2_bench.py
Unverified
1import json, math, sys, time
2import numpy as np
3import torch
4from torch import nn
5
6sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
7from bench import get_dataset, train_model, sweep_baseline, evaluate, make_report
8
9SEEDS = tuple(range(8))
10LR_GRID = [0.0015, 0.003, 0.006]
11EPOCHS = 15
12BATCH = 128
13WIDTH = 32
14RADIUS = 2
15
16
17def dense_attention(h, q, k, v):
18 qq, kk, vv = q(h), k(h), v(h)
19 scores = qq @ kk.transpose(-1, -2) / math.sqrt(qq.shape[-1])
20 return torch.softmax(scores, dim=-1) @ vv
21
22
23def gated_local_global(h, q, k, v, gate, radius=RADIUS):
24 b, n, _ = h.shape
25 qq, kk, vv = q(h), k(h), v(h)
26 offsets = list(range(-radius, radius + 1))
27 positions = torch.arange(n, device=h.device)
28 dst = torch.stack([(positions + off).clamp(0, n - 1) for off in offsets], dim=1)
29 valid = torch.stack([(positions + off >= 0) & (positions + off < n)
30 for off in offsets], dim=1)
31 kn = kk[:, dst]
32 vn = vv[:, dst]
33 scores = (qq[:, :, None, :] * kn).sum(-1) / math.sqrt(qq.shape[-1])
34 scores = scores.masked_fill(~valid[None, :, :], -torch.inf)
35 local = (torch.softmax(scores, dim=-1)[..., None] * vn).sum(dim=2)
36 phi_q = torch.nn.functional.elu(qq) + 1.0
37 phi_k = torch.nn.functional.elu(kk) + 1.0
38 stats = phi_k.transpose(1, 2) @ vv
39 norm = phi_k.sum(dim=1)
40 glob = (phi_q @ stats) / ((phi_q @ norm[:, :, None]).squeeze(-1)[:, :, None] + 1e-6)
41 g = torch.sigmoid(gate(h))
42 return g * local + (1.0 - g) * glob, g
43
44
45class SharedTransformer(nn.Module):
46 def __init__(self, win, out_dim, idea=False):
47 super().__init__()
48 self.idea = idea
49 self.inp = nn.Linear(1, WIDTH)
50 self.pos = nn.Parameter(torch.randn(1, win, WIDTH) * 0.02)
51 self.q = nn.Linear(WIDTH, WIDTH, bias=False)
52 self.k = nn.Linear(WIDTH, WIDTH, bias=False)
53 self.v = nn.Linear(WIDTH, WIDTH, bias=False)
54 self.proj = nn.Linear(WIDTH, WIDTH)
55 self.norm = nn.LayerNorm(WIDTH)
56 self.ff = nn.Sequential(nn.Linear(WIDTH, 64), nn.ReLU(), nn.Linear(64, WIDTH))
57 self.head = nn.Linear(win * WIDTH, out_dim)
58 if idea:
59 self.gate = nn.Linear(WIDTH, 1)
60
61 def forward(self, x, return_gate=False):
62 h = self.inp(x.unsqueeze(-1)) + self.pos[:, :x.shape[1]]
63 if self.idea:
64 att, g = gated_local_global(h, self.q, self.k, self.v, self.gate)
65 else:
66 att = dense_attention(h, self.q, self.k, self.v)
67 g = None
68 h = self.norm(h + self.proj(att))
69 h = h + self.ff(h)
70 out = self.head(h.reshape(h.shape[0], -1))
71 return (out, g) if return_gate else out
72
73
74def run_one(seed, lr, idea, collect=False):
75 torch.manual_seed(seed)
76 np.random.seed(seed)
77 ds = get_dataset("sequence", seed, n_train=400, n_test=400)
78 win = int(ds["xtr"].shape[1])
79 model = SharedTransformer(win, int(ds["ytr"].shape[-1]), idea=idea)
80 trained, metric, history = train_model(model, ds, epochs=EPOCHS, lr=lr,
81 batch=BATCH, weight_decay=0.0, log=lambda *_: None)
82 if trained is None:
83 return float("nan") if not collect else {"metric": float("nan")}
84 if not collect:
85 return float(metric)
86 device = next(trained.parameters()).device
87 x = ds["xte"].to(device)
88 with torch.no_grad():
89 if idea:
90 pred, gates = trained(x, return_gate=True)
91 gate_mean = float(gates.mean().cpu())
92 gate_std = float(gates.std().cpu())
93 else:
94 pred = trained(x)
95 gate_mean = None
96 gate_std = None
97 # Measured forward latency on the trained model, synchronized when CUDA is used.
98 trained.eval()
99 for _ in range(3):
100 with torch.no_grad(): trained(x[:128])
101 if device.type == "cuda": torch.cuda.synchronize(device)
102 t0 = time.perf_counter()
103 for _ in range(10):
104 with torch.no_grad(): trained(x[:128])
105 if device.type == "cuda": torch.cuda.synchronize(device)
106 latency_ms = (time.perf_counter() - t0) * 1000.0 / 10.0
107 return {"metric": float(metric), "gate_mean": gate_mean, "gate_std": gate_std,
108 "latency_ms": latency_ms, "n_nodes": win, "edges_per_node": 2 * RADIUS + 1}
109
110
111def baseline_factory(cfg):
112 return lambda seed: run_one(seed, float(cfg["lr"]), False)
113
114
115def main():
116 # Baseline sweep uses exactly the union of all idea learning rates.
117 grid = [{"lr": lr} for lr in LR_GRID]
118 base = sweep_baseline(baseline_factory, grid, seeds=SEEDS)
119 best_lr = float(base["best_cfg"]["lr"])
120 idea_cfgs = [{"lr": lr} for lr in LR_GRID]
121 idea_runs = {str(c["lr"]): evaluate(lambda seed, lr=c["lr"]: run_one(seed, lr, True), seeds=SEEDS)
122 for c in idea_cfgs}
123 best_key = min(idea_runs, key=lambda k: idea_runs[k]["mean"])
124 idea = idea_runs[best_key]
125 sig = [run_one(s, float(best_key), True, collect=True) for s in SEEDS]
126 base_sig = [run_one(s, best_lr, False, collect=True) for s in SEEDS]
127 signature = {
128 "prediction": "bounded-degree local-global attention scales with O(E*r + N*r^2), while dense attention scales with O(N^2*r)",
129 "observed": {
130 "sequence_length": int(np.mean([z["n_nodes"] for z in sig])),
131 "local_edges_per_node": int(np.mean([z["edges_per_node"] for z in sig])),
132 "idea_latency_ms_mean": float(np.mean([z["latency_ms"] for z in sig])),
133 "baseline_latency_ms_mean": float(np.mean([z["latency_ms"] for z in base_sig])),
134 "idea_gate_mean": float(np.mean([z["gate_mean"] for z in sig])),
135 "idea_gate_std_mean": float(np.mean([z["gate_std"] for z in sig]))
136 },
137 "confirmed": False
138 }
139 # At this fixed short sequence length, observed latency confirms the
140 # mechanism only if the hybrid is faster; otherwise report honestly.
141 signature["confirmed"] = signature["observed"]["idea_latency_ms_mean"] < signature["observed"]["baseline_latency_ms_mean"]
142 report = make_report("sequence", "transformer_tiny", base, idea, signature)
143 report["idea_sweep"] = {k: v for k, v in idea_runs.items()}
144 report["protocol"] = {"seeds": list(SEEDS), "epochs": EPOCHS, "batch": BATCH,
145 "lr_union": LR_GRID, "best_idea_lr": float(best_key)}
146 with open("bench_report.json", "w") as f:
147 json.dump(report, f, indent=2)
148 print(json.dumps(report, indent=2))
149
150
151if __name__ == "__main__":
152 main()