Resolving Landmark Bottleneck / run_experiment.py

✓✓ Beats tuned baseline

Raw ⬇ ZIP
  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()