Critical Cross-Layer Weight Sharing / critical_sharing_experiment.py

Mechanism confirmed, baseline not beaten

Raw ⬇ ZIP
  1import json, math, os
  2import numpy as np
  3
  4SEED = 2716
  5rng = np.random.default_rng(SEED)
  6
  7
  8def corr_matrix(L, gamma):
  9    idx = np.arange(L)
 10    return (1.0 + np.abs(idx[:, None] - idx[None, :])) ** (-gamma)
 11
 12
 13def sums(L, gamma):
 14    t = np.arange(1, L + 1, dtype=float)
 15    c = t ** (-gamma)
 16    return float(np.sum(c)), float(np.sum(c * c))
 17
 18
 19def log_slope(xs, ys):
 20    return float(np.polyfit(np.log(xs), np.log(np.maximum(ys, 1e-30)), 1)[0])
 21
 22
 23def empirical_gaussian_check(L, gamma, n=12000, block=512):
 24    # For x ~ N(0,C), E[(sum x_l^2)^2] = tr(C)^2 + 2 tr(C^2).
 25    C = corr_matrix(L, gamma)
 26    chol = np.linalg.cholesky(C + 1e-10 * np.eye(L))
 27    vals = []
 28    for start in range(0, n, block):
 29        z = rng.normal(size=(min(block, n-start), L))
 30        x = z @ chol.T
 31        vals.append(np.sum(x*x, axis=1))
 32    q = np.concatenate(vals)
 33    empirical = float(np.mean(q*q))
 34    exact = float(np.trace(C)**2 + 2*np.trace(C@C))
 35    # Excess over independent diagonal baseline, divided by L, is governed by sum c_t^2.
 36    ind = float(L*L + 2*L)
 37    return empirical, exact, float((exact-ind)/(2*L))
 38
 39
 40def residual_jacobian_stats(L, gamma, width=32, trials=20, alpha=.8):
 41    # Small fixed residual MLP Jacobian: x_{l+1}=x_l + scale*tanh(W_l x_l).
 42    # W_l has prescribed cross-depth covariance coordinatewise.
 43    d = width
 44    scale = 0.35 / math.sqrt(L)
 45    C = corr_matrix(L, gamma)
 46    chol = np.linalg.cholesky(C + 1e-10*np.eye(L))
 47    singular = []
 48    grad4 = []
 49    for _ in range(trials):
 50        E = rng.normal(size=(L,d,d)) / math.sqrt(d)
 51        G = np.empty((L,d,d))
 52        for i in range(d):
 53            for j in range(d):
 54                G[:,i,j] = chol @ rng.normal(size=L) / math.sqrt(d)
 55        W = math.sqrt(alpha)*G + math.sqrt(1-alpha)*E
 56        J = np.eye(d)
 57        x = rng.normal(size=d)
 58        for l in range(L):
 59            a = W[l] @ x
 60            deriv = 1.0 - np.tanh(a)**2
 61            J = (np.eye(d) + scale*(deriv[:,None]*W[l])) @ J
 62            x = x + scale*np.tanh(a)
 63        sv = np.linalg.svd(J, compute_uv=False)
 64        singular.append([float(np.median(sv)), float(np.max(sv))])
 65        # row-wise Jacobian sensitivity proxy
 66        grad4.append(float(np.mean(np.sum(J*J,axis=1)**2) / (np.mean(np.sum(J*J,axis=1))**2)))
 67    return {'median_sv': float(np.mean(np.array(singular)[:,0])),
 68            'max_sv': float(np.mean(np.array(singular)[:,1])),
 69            'normalized_gradient_fourth': float(np.mean(grad4))}
 70
 71
 72def main():
 73    depths = np.array([32,64,128,256,512,1024,2048,4096,8192,16384,32768])
 74    gammas = [0.3, 0.5, 0.75, 1.0, 1.25]
 75    rows=[]
 76    for g in gammas:
 77        s1=np.array([sums(int(L),g)[0] for L in depths])
 78        s2=np.array([sums(int(L),g)[1] for L in depths])
 79        # Predicted asymptotic slopes: max(1-g,0), max(1-2g,0), with log boundary.
 80        expected_s1 = max(1-g,0.0)
 81        expected_s2 = max(1-2*g,0.0)
 82        empirical, exact, excess = empirical_gaussian_check(128,g)
 83        rows.append({'gamma':g, 's1_slope_observed':log_slope(depths,s1),
 84                     's1_slope_predicted':expected_s1,
 85                     's2_slope_observed':log_slope(depths,s2),
 86                     's2_slope_predicted':expected_s2,
 87                     's1_at_final':float(s1[-1]), 's2_at_final':float(s2[-1]),
 88                     's1_tail_slope_observed':log_slope(depths[-5:],s1[-5:]),
 89                     's2_tail_slope_observed':log_slope(depths[-5:],s2[-5:]),
 90                     's2_normalized_final':float(s2[-1]/(depths[-1]**max(1-2*g,0))) if g<.5 else None,
 91                     'gaussian_q2_empirical_L128':empirical,
 92                     'gaussian_q2_exact_L128':exact,
 93                     'gaussian_excess_per_2L':excess})
 94    # Check critical logarithms by fitting against log L; ratio should stabilize.
 95    critical=[]
 96    for g in [.5,1.0]:
 97        ss=np.array([sums(int(L),g)[1 if g==.5 else 0] for L in depths])
 98        critical.append({'gamma':g,'quantity':'S2' if g==.5 else 'S1',
 99                         'log_fit_slope':log_slope(depths,ss),
100                         'endpoints': [float(ss[0]),float(ss[-1])],
101                         'log_ratio_start':float(ss[0]/math.log(depths[0])),
102                         'log_ratio_end':float(ss[-1]/math.log(depths[-1]))})
103    # Network sanity comparison at moderate depth; initialization is the only changed factor.
104    net=[]
105    for L in [16,32,64,128]:
106        iid = residual_jacobian_stats(L, 100.0, trials=8, alpha=0.0)
107        corr = residual_jacobian_stats(L, .3, trials=8, alpha=1.0)
108        net.append({'depth':L,'iid':iid,'powerlaw_gamma_0.3':corr})
109    out={'seed':SEED,'depths':depths.tolist(),'formula_sweeps':rows,
110         'critical_log_checks':critical,'network_jacobian_sanity':net,
111         'notes':'S1=sum_{t=1}^L t^(-gamma), S2=sum c_t^2. Slopes use finite-depth log-log regression. Gaussian check uses exact Wick formula.'}
112    with open('results.json','w') as f: json.dump(out,f,indent=2)
113    print(json.dumps(out,indent=2))
114
115if __name__=='__main__': main()