Spectral pinning of neural modules / spectral_pinning_experiment.py

✓✓ Beats tuned baseline

Raw ⬇ ZIP
  1import json
  2import numpy as np
  3from scipy.linalg import eigh
  4
  5SEED = 2891
  6rng = np.random.default_rng(SEED)
  7
  8
  9def make_graph(n=8):
 10    # Connected, nonuniform weighted graph (ring plus reproducible chords).
 11    A = np.zeros((n, n), dtype=float)
 12    for i in range(n):
 13        w = 0.25 + 0.10 * ((i * 3) % 4)
 14        j = (i + 1) % n
 15        A[i, j] = A[j, i] = w
 16    for i, j, w in [(0, 3, .90), (1, 5, .18), (2, 6, .55), (3, 7, .28), (0, 6, .12)]:
 17        A[i, j] = A[j, i] = w
 18    L = np.diag(A.sum(1)) - A
 19    return A, L
 20
 21
 22def gap(L, pins, p):
 23    P = np.zeros(L.shape[0])
 24    P[list(pins)] = p
 25    return float(eigh(L + np.diag(P), eigvals_only=True, subset_by_index=[0, 0])[0])
 26
 27
 28def greedy_pins(L, budget, p):
 29    selected = []
 30    current = 0.0
 31    gains = []
 32    for _ in range(budget):
 33        candidates = [i for i in range(L.shape[0]) if i not in selected]
 34        vals = [(gap(L, selected + [i], p), i) for i in candidates]
 35        bestval, best = max(vals, key=lambda z: (z[0], -z[1]))
 36        best = int(best)
 37        gains.append(bestval - current)
 38        selected.append(best)
 39        current = bestval
 40    return selected, gains
 41
 42
 43def rate_from_simulation(L, pins, p, c, Gamma, x0):
 44    """Exact eigendecomposition simulation, then fit the late log-energy slope."""
 45    n, d = L.shape[0], Gamma.shape[0]
 46    P = np.zeros(n); P[list(pins)] = p
 47    M = np.kron(L + np.diag(P), Gamma)
 48    vals, vecs = eigh(M)
 49    z0 = vecs.T @ x0.reshape(-1)
 50    # Enough time for the slowest excited mode to dominate.
 51    lam_pos = vals[vals > 1e-12][0]
 52    times = np.linspace(0.0, max(30.0, 12.0 / (c * lam_pos)), 500)
 53    coeff = z0[:, None] * np.exp(-c * vals[:, None] * times[None, :])
 54    xt = vecs @ coeff
 55    V = 0.5 * np.sum(xt * xt, axis=0)
 56    # fit last half, avoiding numerical underflow
 57    take = slice(250, 470)
 58    slope = float(np.polyfit(times[take], np.log(np.maximum(V[take], 1e-300)), 1)[0])
 59    predicted = -2.0 * c * float(vals[0])
 60    return slope, predicted, float(V[0]), float(V[-1])
 61
 62
 63def main():
 64    A, L = make_graph()
 65    n, d = L.shape[0], 3
 66    Gamma = np.array([[1.0, .15, 0.0], [.15, 1.7, .08], [0.0, .08, 2.4]])
 67    gamma_min = float(eigh(Gamma, eigvals_only=True, subset_by_index=[0, 0])[0])
 68    # Kronecker ordering is module-major, each module has d coordinates.
 69    x0 = rng.normal(size=(n, d))
 70    p = 1.0
 71    budget = 2
 72    greedy, greedy_gains = greedy_pins(L, budget, p)
 73    degrees = A.sum(1)
 74    high_degree = [int(i) for i in np.argsort(-degrees)[:budget]]
 75    # Deterministic random baseline plus all one/two node possibilities for context.
 76    random_sets = [[int(i) for i in x] for x in [(0, 1), (2, 7), (1, 6), (3, 5), (0, 4)]]
 77    random_gaps = [gap(L, s, p) for s in random_sets]
 78    gd = gap(L, greedy, p)
 79    hd = gap(L, high_degree, p)
 80
 81    # Prediction 1: rate = -2*c*lambda_min(Lg)*lambda_min(Gamma), across c.
 82    scaling = []
 83    for c in [0.25, 0.5, 1.0, 2.0]:
 84        obs, pred, _, _ = rate_from_simulation(L, greedy, p, c, Gamma, x0)
 85        scaling.append({'c': c, 'observed_slope': obs, 'predicted_slope': pred,
 86                        'ratio_abs_observed_to_predicted': abs(obs / (pred * gamma_min))})
 87
 88    # Prediction 2: stronger pinning raises the grounded gap with diminishing marginal gain.
 89    pin_sweep = []
 90    for strength in [0.05, 0.1, 0.25, 0.5, 1.0, 2.0, 4.0]:
 91        g = gap(L, greedy, strength)
 92        obs, _, _, _ = rate_from_simulation(L, greedy, strength, 1.0, Gamma, x0)
 93        predicted_rate = -2.0 * g * gamma_min
 94        pin_sweep.append({'p': strength, 'lambda_min': g, 'observed_slope': obs,
 95                          'predicted_slope': predicted_rate,
 96                          'ratio': abs(obs / predicted_rate)})
 97
 98    # Prediction 3: spectral greedy produces the largest gap among candidate choices.
 99    comparisons = {'greedy': {'pins': greedy, 'gap': gd},
100                   'high_degree': {'pins': high_degree, 'gap': hd},
101                   'random_sets': [{'pins': s, 'gap': g} for s, g in zip(random_sets, random_gaps)]}
102    for name, pins in [('greedy', greedy), ('high_degree', high_degree)]:
103        obs, pred0, _, _ = rate_from_simulation(L, pins, p, 1.0, Gamma, x0)
104        comparisons[name]['observed_slope'] = obs
105        comparisons[name]['predicted_slope'] = pred0
106    random_slopes = []
107    for s in random_sets:
108        obs, pred0, _, _ = rate_from_simulation(L, s, p, 1.0, Gamma, x0)
109        random_slopes.append(obs)
110    comparisons['random_mean_gap'] = float(np.mean(random_gaps))
111    comparisons['random_mean_observed_slope'] = float(np.mean(random_slopes))
112
113    all_ratios = [x['ratio_abs_observed_to_predicted'] for x in scaling] + [x['ratio'] for x in pin_sweep]
114    output = {
115        'seed': SEED, 'graph_nodes': n, 'module_dimension': d,
116        'lambda_min_Gamma': gamma_min, 'degrees': degrees.tolist(),
117        'predictions': {
118            'decay_rate': 'late log(V) slope = -2*c*lambda_min(L+P)*lambda_min(Gamma)',
119            'pinning': 'lambda_min and decay magnitude increase with p, with diminishing gains',
120            'selection': 'greedy marginal gap maximization beats random choices on this graph'
121        },
122        'greedy_pins': greedy, 'greedy_marginal_gains': greedy_gains,
123        'scaling_sweep': scaling, 'pin_strength_sweep': pin_sweep,
124        'selection_comparison': comparisons,
125        'max_rate_prediction_ratio_error': float(max(abs(r - 1.0) for r in all_ratios)),
126        'all_rate_ratios_within_25pct': bool(all(0.75 <= r <= 1.25 for r in all_ratios)),
127        'greedy_beats_random_mean_gap': bool(gd > np.mean(random_gaps)),
128        'greedy_beats_high_degree_gap': bool(gd > hd),
129    }
130    with open('results.json', 'w') as f:
131        json.dump(output, f, indent=2, default=lambda x: x.item() if hasattr(x, 'item') else x)
132    print(json.dumps(output, indent=2))
133
134
135if __name__ == '__main__':
136    main()