Histogram-Controlled Cluster Updates for Iterative GNNs / experiment.py

Mechanism confirmed, baseline not beaten

Raw ⬇ ZIP
  1import json, math
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 1172
  6rng = np.random.default_rng(SEED)
  7
  8# Histogram mechanism from the proposal.
  9def histogram(omega, Amax, R):
 10    omega = np.asarray(omega, dtype=int)
 11    q = np.minimum(R - 1, np.floor(R * np.clip(omega, 0, Amax) / Amax).astype(int))
 12    return np.bincount(q, minlength=R).astype(float) / len(omega)
 13
 14def quantizer_sweep():
 15    # Prediction: bin r starts at the smallest integer ceil(r*Amax/R),
 16    # except bin zero; clipping makes the final bin absorb the endpoint.
 17    A, R = 11, 4
 18    observed = []
 19    for r in range(1, R):
 20        vals = np.arange(A + 1)
 21        qs = np.minimum(R-1, np.floor(R*vals/A).astype(int))
 22        observed.append(int(vals[qs == r][0]))
 23    predicted = [int(math.ceil(r*A/R)) for r in range(1, R)]
 24    return {"Amax": A, "R": R, "predicted_first_integer": predicted,
 25            "observed_first_integer": observed,
 26            "max_threshold_error": int(max(abs(a-b) for a,b in zip(predicted, observed)))}
 27
 28def invariance_sweep():
 29    # Prediction: any permutation of nodes, and any replicated cluster with
 30    # identical empirical distribution, has exactly the same fixed-size state.
 31    A, R = 8, 5
 32    omega = np.array([0, 1, 2, 2, 5, 8, 3, 3, 1, 0])
 33    h = histogram(omega, A, R)
 34    perm_errors, replicate_errors = [], []
 35    for _ in range(100):
 36        perm_errors.append(float(np.max(np.abs(h - histogram(rng.permutation(omega), A, R)))))
 37        k = int(rng.integers(1, 8))
 38        replicate_errors.append(float(np.max(np.abs(h - histogram(np.tile(omega, k), A, R)))))
 39    # Deliberately show cardinality is not represented: same distribution, sizes differ.
 40    return {"permutation_max_abs_error": max(perm_errors),
 41            "replication_max_abs_error": max(replicate_errors),
 42            "state_dimension": R,
 43            "sizes_tested": [len(omega), len(omega)*7]}
 44
 45def synchronous_order_sweep():
 46    # Two-node linear update x_i <- (1-eta)x_i + eta*x_neighbor.
 47    # For x=[1,0], synchronous vs sequential (0 then 1) differs exactly eta^2
 48    # in the second component, demonstrating why pre-update gathering matters.
 49    etas = np.linspace(.05, .95, 10)
 50    rows = []
 51    for eta in etas:
 52        x = np.array([1., 0.])
 53        sync = np.array([(1-eta)*x[0] + eta*x[1], (1-eta)*x[1] + eta*x[0]])
 54        seq = x.copy()
 55        seq[0] = (1-eta)*seq[0] + eta*seq[1]
 56        seq[1] = (1-eta)*seq[1] + eta*seq[0]
 57        err = float(np.max(np.abs(sync-seq)))
 58        rows.append((float(eta), err, float(eta**2)))
 59    rel = max(abs(e-p) for _,e,p in rows)
 60    return {"rows": rows, "predicted_linf_error": "eta^2",
 61            "max_absolute_error_vs_prediction": rel}
 62
 63def make_problem(n=96, clusters=8, seed=SEED):
 64    r = np.random.default_rng(seed)
 65    # Ring plus random edges; target is a smooth binary-ish signal.
 66    edges = set()
 67    for i in range(n):
 68        for d in (1, 2):
 69            j=(i+d)%n; edges.add((min(i,j),max(i,j)))
 70    for _ in range(2*n):
 71        i,j=r.integers(0,n,2)
 72        if i != j: edges.add((min(i,j),max(i,j)))
 73    adj=[set() for _ in range(n)]
 74    for i,j in edges: adj[i].add(j); adj[j].add(i)
 75    target=np.sin(np.arange(n)*2*np.pi/24)
 76    x=r.normal(0,.9,n)
 77    blocks=[np.arange(a*n//clusters,(a+1)*n//clusters) for a in range(clusters)]
 78    return adj,target,x,blocks
 79
 80def residual_mass(x, target, adj):
 81    # Residual indicator is thresholded local error, as in the proposal.
 82    return float(np.sum(np.abs(x-target) > .15))
 83
 84def update_cluster(x, target, adj, ids, eta=.55):
 85    # All neighbor reads are from x, then scatter: synchronous within cluster.
 86    old=x.copy(); out=[]
 87    for i in ids:
 88        neigh=list(adj[i])
 89        mean=np.mean(old[neigh]) if neigh else old[i]
 90        out.append((1-eta)*old[i] + eta*(.55*mean+.45*target[i]))
 91    x[np.asarray(ids)] = out
 92
 93def mini_compare():
 94    adj,target,x0,blocks=make_problem()
 95    # Standard: one full synchronous sweep is the natural reference.
 96    full=[]; x=x0.copy()
 97    for t in range(30):
 98        old=x.copy()
 99        for ids in blocks:
100            # full sweep uses a single pre-update state for every node
101            update_cluster(x, target, adj, np.arange(len(x)), eta=.55)
102            break
103        full.append(residual_mass(x,target,adj))
104    # Idea: observe per-cluster residual histograms and greedily select the
105    # cluster with greatest high-bin mass; this is the exploitation behavior
106    # a learned Q scheduler is intended to approximate.
107    idea=[]; x=x0.copy(); decisions=0
108    for t in range(30*len(blocks)):
109        scores=[]
110        for ids in blocks:
111            err=np.abs(x[ids]-target[ids]); omega=np.minimum(8,(err>.15).astype(int)*8)
112            scores.append(histogram(omega,8,4)[-1])
113        a=int(np.argmax(scores)); update_cluster(x,target,adj,blocks[a],eta=.55); decisions+=1
114        if decisions % len(blocks)==0: idea.append(residual_mass(x,target,adj))
115    return {"full_synchronous_residual_each_sweep": full,
116            "histogram_selected_residual_each_8_decisions": idea,
117            "final_full": full[-1], "final_idea": idea[-1],
118            "idea_decisions_for_30_sweep_equivalents": decisions,
119            "note":"This is a tiny signal check, not a claim of wall-clock speedup."}
120
121def main():
122    out={"seed":SEED,"predictions":{
123      "P1_histogram_permutation_and_cardinality_invariance": invariance_sweep(),
124      "P2_quantizer_transition_thresholds": quantizer_sweep(),
125      "P3_synchronous_cluster_update_order_error": synchronous_order_sweep()},
126      "mini_experiment":mini_compare()}
127    Path("results.json").write_text(json.dumps(out, indent=2))
128    print(json.dumps(out, indent=2))
129if __name__ == '__main__': main()