Spiderweb Hierarchical Attention / spiderweb_experiment.py
Beats tuned baseline
1"""MVP verification for Spiderweb Hierarchical Attention.
2Run: python3 spiderweb_experiment.py
3"""
4import math, json, random
5import numpy as np
6
7
8def hyperbolic_distance(dx, y, C=1.0):
9 # Half-space metric with equal heights: C arcosh(1+dx^2/(2 y^2)).
10 z = 1.0 + (dx * dx) / (2.0 * y * y)
11 return C * np.arccosh(np.maximum(z, 1.0))
12
13
14def preferred_level(dx, n, L, eta, y0, tau):
15 levels = np.arange(L + 1)
16 ys = y0 * eta ** levels
17 ds = hyperbolic_distance(dx, ys)
18 return int(np.argmin(np.abs(ds - tau))), ds
19
20
21def predicted_transition(dx, L, eta, y0, tau):
22 # Continuous solution of d(dx,y)=tau, then nearest/clipped dyadic level.
23 ystar = dx / math.sqrt(2.0 * (math.cosh(tau) - 1.0))
24 raw = math.log(ystar / y0, eta)
25 return min(L, max(0, int(round(raw)))), raw
26
27
28def hierarchy_work(n, neighbor_radius=1):
29 # One summary and one fixed-radius horizontal attention neighborhood per level;
30 # token-level gates/broadcasts are counted, as required by the pseudocode.
31 cells = n
32 horizontal = 0
33 broadcasts = 0
34 levels = 0
35 while cells >= 1:
36 horizontal += cells * (min(cells - 1, neighbor_radius) * 2 + 1)
37 broadcasts += n
38 levels += 1
39 if cells == 1: break
40 cells = (cells + 1) // 2
41 return horizontal + broadcasts, levels, horizontal, broadcasts
42
43
44def graph_coverage(n, kind, window=2, radius=1):
45 """Reachable nodes after one communication layer; boolean adjacency."""
46 A = np.zeros((n, n), dtype=bool)
47 if kind == 'dense':
48 A[:] = True
49 elif kind == 'local':
50 for i in range(n):
51 A[i, max(0, i-window):min(n, i+window+1)] = True
52 elif kind == 'spiderweb':
53 # Undirected approximation to pooling, horizontal neighboring cells,
54 # and broadcast through every dyadic level.
55 L = int(math.ceil(math.log2(n)))
56 for i in range(n):
57 for l in range(L + 1):
58 size = 2 ** l
59 lo = (i // size) * size; hi = min(n, lo + size)
60 A[i, lo:hi] = True
61 for off in range(1, radius + 1):
62 for c in (lo - off*size, hi + (off-1)*size):
63 if 0 <= c < n:
64 A[i, c:min(n, c+size)] = True
65 return A
66
67
68def main():
69 np.random.seed(7); random.seed(7)
70 eta, y0, tau, n, L = 2.0, 0.02, 1.0, 256, 8
71 # Prediction 1: d(dx,y_l) strictly decreases with level for every dx>0.
72 dxs = np.array([1/n, 8/n, 32/n, 100/n])
73 monotone = []
74 for dx in dxs:
75 _, ds = preferred_level(dx, n, L, eta, y0, tau)
76 monotone.append(bool(np.all(np.diff(ds) < 0)))
77 # Prediction 2: scale transition agrees with closed form inverse metric.
78 rows = []
79 for dx in np.linspace(1/n, 0.95, 25):
80 obs, ds = preferred_level(float(dx), n, L, eta, y0, tau)
81 pred, raw = predicted_transition(float(dx), L, eta, y0, tau)
82 rows.append((float(dx), obs, pred, abs(obs-pred)))
83 errors = np.array([r[3] for r in rows])
84 # Prediction 3: work/(n log2 n) stays bounded while dense/(n log n) grows.
85 ns = np.array([32,64,128,256,512,1024,2048])
86 sw = np.array([hierarchy_work(int(x))[0] for x in ns])
87 dense = ns.astype(float)**2
88 norm_sw = sw/(ns*np.log2(ns)); norm_dense = dense/(ns*np.log2(ns))
89 # Toy coverage and operations at n=256.
90 cov = {}
91 for kind in ('dense','local','spiderweb'):
92 A = graph_coverage(n, kind)
93 cov[kind] = {'reachable_from_middle': int(A[n//2].sum()),
94 'mean_reachable': float(A.sum()/n)}
95 result = {
96 'settings': {'n':n,'eta':eta,'y0':y0,'tau':tau,'levels':L},
97 'prediction_1_monotone_distance': {'dx':dxs.tolist(),'confirmed_each':monotone,
98 'claim':'equal-height hyperbolic distance strictly decreases as y_l grows'},
99 'prediction_2_transition': {'formula':'y*=dx/sqrt(2(cosh(tau)-1)); l*=round(log_eta(y*/y0))',
100 'mean_abs_level_error':float(errors.mean()), 'max_abs_level_error':float(errors.max()),
101 'fraction_exact':float(np.mean(errors==0)), 'samples':rows[::6]},
102 'prediction_3_work_scaling': {'n':ns.tolist(),'work':sw.tolist(),
103 'work_over_n_log2n':norm_sw.tolist(),'dense_over_n_log2n':norm_dense.tolist(),
104 'spiderweb_ratio_range':[float(norm_sw.min()),float(norm_sw.max())],
105 'dense_ratio_growth':float(norm_dense[-1]/norm_dense[0])},
106 'toy_communication': {'spiderweb_work':hierarchy_work(n)[0],
107 'dense_attention_scores':n*n,'local_attention_scores':n*5,'coverage':cov}
108 }
109 print(json.dumps(result, indent=2))
110
111if __name__ == '__main__': main()
112
113class SpiderwebAttention:
114 """Small readable NumPy implementation of the proposed hierarchy.
115
116 This is an inference/MVP layer: cell pooling is mean pooling, horizontal
117 attention uses nearby cells, and each level broadcasts a gated summary.
118 """
119 def __init__(self, dim, eta=2, levels=None, radius=1, seed=0):
120 self.dim, self.eta, self.radius = dim, int(eta), int(radius)
121 self.rng = np.random.default_rng(seed)
122 self.levels = levels
123 self.gates = None
124
125 def __call__(self, x):
126 x = np.asarray(x, dtype=float)
127 n, d = x.shape
128 L = self.levels if self.levels is not None else int(math.ceil(math.log2(n)))
129 cur = x.copy()
130 outputs = []
131 for lev in range(L + 1):
132 size = self.eta ** lev
133 cells = [(a, min(a + size, n)) for a in range(0, n, size)]
134 s = np.stack([cur[a:b].mean(0) for a, b in cells])
135 q = s / math.sqrt(d)
136 scores = q @ s.T
137 mask = np.full(scores.shape, -np.inf)
138 for c in range(len(cells)):
139 lo, hi = max(0, c-self.radius), min(len(cells), c+self.radius+1)
140 mask[c, lo:hi] = scores[c, lo:hi]
141 mask -= np.max(mask, axis=1, keepdims=True)
142 att = np.exp(mask); att /= att.sum(axis=1, keepdims=True)
143 updated = att @ s
144 broadcast = np.zeros_like(cur)
145 for c, (a, b) in enumerate(cells): broadcast[a:b] = updated[c]
146 gate = 1.0 / (lev + 1.0)
147 cur = cur + gate * broadcast
148 outputs.append(cur.copy())
149 return cur, outputs
150
151
152def dense_mean(x):
153 return x + np.broadcast_to(x.mean(0), x.shape)
154
155
156def local_mean(x, window=2):
157 out = x.copy()
158 for i in range(len(x)):
159 out[i] += x[max(0, i-window):min(len(x), i+window+1)].mean(0)
160 return out