Rank-Normalized Nonlinear Spectral Preconditioner / rank_spectral_mvp.py
Failed on benchmark
1import json, math, time
2import numpy as np
3from scipy.special import ndtri
4
5SEED = 1163
6rng = np.random.default_rng(SEED)
7
8def rank_normalize(X):
9 n, p = X.shape
10 Z = np.empty_like(X, dtype=float)
11 for j in range(p):
12 order = np.argsort(X[:, j], kind='mergesort')
13 probs = (np.arange(n) + 0.5) / n
14 zcol = ndtri(probs)
15 Z[order, j] = zcol
16 Z -= Z.mean(0, keepdims=True)
17 return Z
18
19def cov(X):
20 X = X - X.mean(0, keepdims=True)
21 return X.T @ X / len(X)
22
23def spectral_clean(S, preserve_top=1):
24 # MVP nonlinear spectral map: retain isolated spikes, pull noisy bulk to
25 # its median. This is a monotone eigenvalue-specific pilot for g_H,gamma.
26 w, V = np.linalg.eigh(S)
27 w = np.maximum(w, 1e-8)
28 p = len(w)
29 bulk = w[:-preserve_top] if preserve_top and p > preserve_top else w
30 med = np.median(bulk)
31 # A mild nonlinear map: bulk is contracted toward its robust center;
32 # top isolated directions are retained, not globally rescaled.
33 cleaned = w.copy()
34 if len(bulk):
35 lo, hi = np.quantile(bulk, [0.1, 0.9])
36 mask = np.arange(p) < p-preserve_top if preserve_top else np.ones(p, bool)
37 cleaned[mask] = med + 0.35 * (w[mask] - med)
38 cleaned[mask] = np.maximum(cleaned[mask], 0.05 * med)
39 return (V * cleaned) @ V.T, w, cleaned
40
41def latent_cov(p, spike):
42 Q, _ = np.linalg.qr(rng.normal(size=(p, p)))
43 vals = np.ones(p); vals[0] = spike
44 return (Q * vals) @ Q.T, Q, vals
45
46def experiment_copula_convergence():
47 # Monotone heavy-tailed marginals preserve ranks but distort Pearson covariance.
48 p, rho, reps = 8, 0.55, 30
49 R = np.full((p,p), rho); np.fill_diagonal(R, 1.)
50 rows=[]
51 for n in [64, 128, 256, 512, 1024]:
52 raw_err=[]; rank_err=[]
53 for _ in range(reps):
54 G = rng.multivariate_normal(np.zeros(p), R, size=n)
55 # strongly nonlinear monotone marginals, finite but very heavy-tailed
56 X = np.sign(G) * (np.exp(np.minimum(np.abs(G), 5.0)) - 1.0)
57 raw_err.append(np.linalg.norm(cov(X)-R, 'fro')/np.linalg.norm(R,'fro'))
58 rank_err.append(np.linalg.norm(cov(rank_normalize(X))-R, 'fro')/np.linalg.norm(R,'fro'))
59 rows.append({'n':n,'raw_relerr':float(np.mean(raw_err)), 'rank_relerr':float(np.mean(rank_err)),
60 'rank_sd':float(np.std(rank_err)/math.sqrt(reps))})
61 # predicted sampling law: rank error ~ n^-1/2
62 slope = np.polyfit(np.log([x['n'] for x in rows]), np.log([x['rank_relerr'] for x in rows]), 1)[0]
63 return rows, float(slope)
64
65def experiment_outliers():
66 # One contaminated fraction with arbitrarily large positive shifts.
67 p,n,reps,q = 8,256,40,0.05
68 R=np.full((p,p),.35); np.fill_diagonal(R,1.)
69 rows=[]
70 for A in [0, 2, 5, 10, 20, 50]:
71 raw=[]; rank=[]; raw_var=[]; rank_var=[]
72 for _ in range(reps):
73 G=rng.multivariate_normal(np.zeros(p),R,size=n)
74 X=G.copy()
75 m=max(1,int(q*n)); ix=rng.choice(n,m,replace=False)
76 X[ix,0] += A
77 raw.append(np.linalg.norm(cov(X)-R,'fro')/np.linalg.norm(R,'fro'))
78 Z=rank_normalize(X)
79 rank.append(np.linalg.norm(cov(Z)-R,'fro')/np.linalg.norm(R,'fro'))
80 raw_var.append(cov(X)[0,0]); rank_var.append(cov(Z)[0,0])
81 rows.append({'A':A,'raw_relerr':float(np.mean(raw)),'rank_relerr':float(np.mean(rank)),
82 'raw_var0':float(np.mean(raw_var)),'rank_var0':float(np.mean(rank_var))})
83 # predicted raw variance asymptote is q(1-q) A^2; fit coefficient at large A.
84 As=np.array([r['A'] for r in rows[-3:]],float)
85 vs=np.array([r['raw_var0'] for r in rows[-3:]])
86 coef=float(np.polyfit(As**2,vs,1)[0])
87 return rows, coef, q*(1-q)
88
89def experiment_spike():
90 # Rank scores plus eigenvalue-specific pilot cleaning: bulk spread shrinks,
91 # while a separated top direction is intentionally retained.
92 n,p,reps=128,12,30
93 rows=[]
94 for spike in [1,2,4,8]:
95 rawe=[]; cle=[]; topraw=[]; topclean=[]
96 C,Q,vals=latent_cov(p,spike)
97 for _ in range(reps):
98 G=rng.multivariate_normal(np.zeros(p), C, size=n)
99 # marginal nonlinearities make Pearson S especially misleading
100 X=np.sign(G)*(np.exp(np.minimum(np.abs(G),4))-1)
101 S=cov(rank_normalize(X))
102 Cc,w,wc=spectral_clean(S,1)
103 rawe.append(np.linalg.norm(S-C,'fro'))
104 cle.append(np.linalg.norm(Cc-C,'fro'))
105 topraw.append(w[-1]); topclean.append(wc[-1])
106 rows.append({'spike':spike,'raw_cov_error':float(np.mean(rawe)),
107 'clean_cov_error':float(np.mean(cle)), 'top_raw':float(np.mean(topraw)),
108 'top_clean':float(np.mean(topclean)), 'clean_win':float(np.mean(cle)<np.mean(rawe))})
109 return rows
110
111
112
113def invsqrt_psd(S, eps=1e-3):
114 w, V = np.linalg.eigh(S)
115 scale = max(float(np.mean(w)), 1e-8)
116 w = np.maximum(w, eps * scale)
117 return (V * (1.0 / np.sqrt(w))) @ V.T
118
119def preprocess_batch(X, mode):
120 if mode == 'raw':
121 return X
122 if mode == 'diag':
123 d = np.sqrt(np.mean((X-X.mean(0))**2, axis=0) + 1e-3)
124 return (X-X.mean(0)) / d
125 Z = rank_normalize(X)
126 S = cov(Z)
127 Cclean, _, _ = spectral_clean(S, preserve_top=1)
128 # Use the specified cleaned covariance to whiten activations.
129 return (X-X.mean(0)) @ invsqrt_psd(Cclean, eps=1e-2)
130
131def experiment_optimization():
132 # Online linear regression with heavy-tailed, contaminated activations.
133 # Same pre-generated batches are used by every method.
134 p, batch, steps, reps = 16, 64, 350, 8
135 lr = 0.025
136 checkpoints = [50, 100, 200, 350]
137 allres = {m: [] for m in ['raw','diag','rank_spectral']}
138 for rep in range(reps):
139 local = np.random.default_rng(SEED + 1000 + rep)
140 beta = local.normal(size=p); beta /= np.linalg.norm(beta)
141 batches=[]
142 for _ in range(steps):
143 G = local.normal(size=(batch,p))
144 # Correlated latent directions, monotone heavy tails, plus rare shifts.
145 G[:,1] = .65*G[:,0] + .76*G[:,1]
146 X = np.sign(G) * (np.exp(np.minimum(np.abs(G), 4.0))-1.0)
147 bad = local.random(batch) < .05
148 X[bad,0] += local.choice([-1,1], size=bad.sum()) * 15.0
149 y = X @ beta + .05*local.normal(size=batch)
150 batches.append((X,y))
151 for mode in allres:
152 w=np.zeros(p); losses=[]
153 for step,(X,y) in enumerate(batches,1):
154 Xt=preprocess_batch(X, mode)
155 pred=Xt@w
156 err=pred-y
157 losses.append(float(np.mean(err*err)))
158 w -= lr * (Xt.T@err)/batch
159 if step in checkpoints:
160 allres[mode].append((step, losses[-1]))
161 summary={}
162 for mode, vals in allres.items():
163 summary[mode]={str(k): float(np.mean([v for s,v in vals if s==k])) for k in checkpoints}
164 return summary
165
166# Preserve the original verification output and add the optimization result.
167if __name__ == '__main__':
168 t=time.time()
169 conv,slope=experiment_copula_convergence()
170 out,coef,pred=experiment_outliers()
171 spike=experiment_spike()
172 opt=experiment_optimization()
173 result={'seed':SEED,'predictions':{
174 'P1_rank_sampling_scaling':'rank covariance error should scale as n^-1/2 after monotone marginal distortion',
175 'P2_outlier_influence':'raw variance coefficient versus A^2 should approach q(1-q)=%.4f; rank variance should remain bounded'%pred,
176 'P3_eigen_specific_cleaning':'bulk contraction should reduce covariance error while retaining top spike'},
177 'observed':{'P1_rows':conv,'P1_loglog_slope':slope,'P2_rows':out,'P2_raw_A2_coefficient':coef,'P2_predicted_coefficient':pred,'P3_rows':spike,'optimization_mse':opt},
178 'elapsed_sec':time.time()-t}
179 print(json.dumps(result, indent=2))