Resolving Landmark Bottleneck / run_experiment.py
Beats tuned baseline
1import json, math, time
2from pathlib import Path
3import numpy as np
4import torch
5from torch import nn
6
7SEED = 421
8np.random.seed(SEED)
9torch.manual_seed(SEED)
10
11
12def make_graph(n=600, k=3, p_in=.16, p_out=.035):
13 labels = np.arange(n) % k
14 rng = np.random.default_rng(SEED)
15 A = np.zeros((n,n), dtype=np.uint8) # A[u,v] means u -> v
16 for u in range(n):
17 same = labels == labels[u]
18 probs = np.where(same, p_in, p_out)
19 row = rng.random(n) < probs
20 row[u] = False
21 A[u] = row
22 # ensure all vertices have a modest amount of incoming structure
23 return A, labels
24
25
26def signatures(A, landmarks):
27 return A[np.asarray(landmarks), :].T.astype(np.float32)
28
29
30def collision_count(Z):
31 return int(Z.shape[0] - np.unique(Z, axis=0).shape[0])
32
33
34def pairwise_c(A):
35 # exact minimum symmetric difference of incoming neighborhoods
36 B = A.T.astype(np.uint8)
37 c = n_pairs = 10**9
38 for i in range(len(B)):
39 d = np.sum(np.bitwise_xor(B[i+1:], B[i]), axis=1) if i+1 < len(B) else np.array([])
40 if len(d): c = min(c, int(d.min()))
41 return c
42
43
44def greedy_landmarks(A, s):
45 n = len(A)
46 incoming = A.T.astype(bool)
47 # distinguish all pairs; represent only currently colliding pairs
48 Z = np.zeros((n, 0), dtype=np.uint8)
49 chosen = []
50 remaining = set(range(n))
51 for _ in range(min(s,n)):
52 if collision_count(Z) == 0: break
53 # Candidate score: number of currently equal-signature pairs split by q.
54 groups = {}
55 for i, row in enumerate(Z): groups.setdefault(row.tobytes(), []).append(i)
56 best, bestscore = None, -1
57 for q in remaining:
58 score = 0
59 bit = incoming[:,q]
60 for inds in groups.values():
61 if len(inds) > 1:
62 score += int(np.sum(bit[inds]) * (len(inds)-np.sum(bit[inds])))
63 if score > bestscore: best, bestscore = q, score
64 if best is None or bestscore <= 0: break
65 chosen.append(best); remaining.remove(best)
66 Z = np.column_stack([Z, incoming[:,best].astype(np.uint8)])
67 return chosen
68
69
70class MLP(nn.Module):
71 def __init__(self, d, hidden=48, k=3):
72 super().__init__(); self.net=nn.Sequential(nn.Linear(d,hidden),nn.ReLU(),nn.Linear(hidden,k))
73 def forward(self,x): return self.net(x)
74
75
76def train_eval(X, y, epochs=100):
77 # Same fixed classifier for every representation, with fixed split and seed.
78 torch.manual_seed(SEED)
79 n=len(y); perm=np.random.default_rng(SEED).permutation(n)
80 tr=perm[:int(.7*n)]; te=perm[int(.7*n):]
81 model=MLP(X.shape[1], k=int(y.max()+1))
82 opt=torch.optim.Adam(model.parameters(), lr=.015, weight_decay=1e-4)
83 xt=torch.tensor(X); yt=torch.tensor(y, dtype=torch.long)
84 t=time.perf_counter()
85 for _ in range(epochs):
86 opt.zero_grad(); loss=nn.functional.cross_entropy(model(xt[tr]),yt[tr]); loss.backward(); opt.step()
87 with torch.no_grad():
88 pred=model(xt[te]).argmax(1); acc=float((pred==yt[te]).float().mean())
89 return acc, float(loss), time.perf_counter()-t
90
91
92def math_check(A, s, draws=4000):
93 n=len(A); incoming=A.T.astype(bool)
94 rng=np.random.default_rng(SEED+1)
95 # Select a representative pair and also evaluate all-pair aggregate prediction.
96 pairs=[(0,1),(0,n//2),(1,n//2)]
97 empirical=[]; exact=[]
98 for u,v in pairs:
99 diff=int(np.sum(np.logical_xor(incoming[u],incoming[v])))
100 p=math.comb(n-diff,s)/math.comb(n,s) if n-diff>=s else 0.0
101 hit=0
102 for _ in range(draws):
103 S=rng.choice(n,s,replace=False)
104 if not np.any(incoming[u,S] != incoming[v,S]): hit+=1
105 empirical.append(hit/draws); exact.append(p)
106 # Aggregate exact bound and Monte Carlo mean collision count.
107 c=pairwise_c(A)
108 bound=math.comb(n,2)*(math.comb(n-c,s)/math.comb(n,s) if n-c>=s else 0.0)
109 obs=[]
110 for _ in range(200):
111 S=rng.choice(n,s,replace=False); obs.append(collision_count(incoming[:,S].astype(np.uint8)))
112 return {'n':n,'s':s,'c_exact':c,'pair_exact':exact,'pair_empirical':empirical,
113 'expected_collision_upper_bound':bound,'sampled_mean_collisions':float(np.mean(obs)),
114 'max_pair_abs_error':float(max(abs(a-b) for a,b in zip(empirical,exact)))}
115
116
117def main():
118 A,y=make_graph(); n=len(y)
119 s=max(4, int(math.ceil(2*math.sqrt(n))))
120 check=math_check(A,s)
121 rng=np.random.default_rng(SEED+2)
122 deg=A.sum(axis=1)+A.sum(axis=0)
123 methods={
124 'uniform':rng.choice(n,s,replace=False).tolist(),
125 'degree':np.argsort(-deg)[:s].tolist(),
126 'greedy':greedy_landmarks(A,s)
127 }
128 # noisy features are deliberately insufficient alone; signatures carry block structure.
129 X0=rng.normal(0,1,(n,8)).astype(np.float32)
130 results={}
131 for name,S in methods.items():
132 t=time.perf_counter(); Z=signatures(A,S); prep=time.perf_counter()-t
133 # x + landmark binary positional signature, exactly the proposed interface
134 acc,loss,train=train_eval(np.concatenate([X0,Z],axis=1),y)
135 results[name]={'s':len(S),'collisions':collision_count(Z),'collision_rate':collision_count(Z)/n,
136 'accuracy':acc,'final_train_loss':loss,'preprocess_sec':prep,'train_sec':train}
137 # Standard control: one-step full incoming-neighbor mean as graph-to-token interface.
138 t=time.perf_counter(); Xfull=np.concatenate([X0, A.T.astype(np.float32)@X0/max(1,A.sum(axis=0).max())],axis=1); prep=time.perf_counter()-t
139 acc,loss,train=train_eval(Xfull,y)
140 results['full_1hop']={'features':Xfull.shape[1],'collisions':'n/a','accuracy':acc,'final_train_loss':loss,'preprocess_sec':prep,'train_sec':train}
141 out={'seed':SEED,'graph':{'n':n,'edges':int(A.sum()),'classes':3},'math_check':check,'methods':results}
142 Path('results.json').write_text(json.dumps(out,indent=2))
143 print(json.dumps(out,indent=2))
144
145if __name__=='__main__': main()