Critical Cross-Layer Weight Sharing / critical_sharing_experiment.py
Mechanism confirmed, baseline not beaten
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()