Anchored Whitening Layer / experiment.py
Unverified
1import json, math
2import numpy as np
3from scipy.linalg import eigh
4from sklearn.linear_model import LogisticRegression
5from sklearn.model_selection import train_test_split
6from sklearn.metrics import accuracy_score
7
8SEED = 2340
9rng = np.random.default_rng(SEED)
10
11
12def equicorr(p, c):
13 return (1-c)*np.eye(p) + c*np.ones((p,p))
14
15
16def invsqrt(R, eps=1e-10):
17 w, U = eigh(R)
18 return (U * (1.0/np.sqrt(np.maximum(w, eps))) ) @ U.T
19
20
21def sqrtm(R):
22 w, U = eigh(R)
23 return (U * np.sqrt(np.maximum(w, 0))) @ U.T
24
25
26def anchored_whitener(R, Q=None, eps=1e-10):
27 if Q is None: Q = np.eye(R.shape[0])
28 return invsqrt(R, eps) @ Q
29
30
31def metrics(R, T):
32 C = T.T @ R @ T
33 fidelity = np.diag(R @ T)
34 return float(np.max(np.abs(C-np.eye(R.shape[0])))), fidelity, float(np.linalg.norm(C-np.eye(R.shape[0]), 'fro'))
35
36
37def analytic_identity_fidelity(p, c):
38 # diag(sqrt(R)) for equicorrelation R
39 return (math.sqrt(1+(p-1)*c) + (p-1)*math.sqrt(1-c))/p
40
41
42def run_math_sweeps():
43 p = 6
44 # Prediction 1: exact covariance residual is numerical precision and does not depend on c.
45 decor = []
46 for c in [0.0, .1, .3, .6, .9]:
47 R = equicorr(p,c); T = anchored_whitener(R)
48 d, f, fn = metrics(R,T)
49 decor.append({'c':c, 'max_cov_error':d, 'fidelity_min':float(f.min())})
50
51 # Prediction 2: Q=I fidelity follows the closed form, measured over correlation sweep.
52 fidelity = []
53 for c in [0.0, .1, .3, .6, .9]:
54 R = equicorr(p,c); T = anchored_whitener(R)
55 observed = float(metrics(R,T)[1].min())
56 predicted = analytic_identity_fidelity(p,c)
57 fidelity.append({'c':c, 'predicted':predicted, 'observed':observed,
58 'abs_error':abs(predicted-observed)})
59
60 # Prediction 3: for an identity-initialized anchored layer, the declared threshold
61 # is satisfied exactly up to rho0(c)=diag(sqrt(R)); above it, the initial anchor violates it.
62 transition=[]
63 for c in [0.0, .1, .3, .6, .9]:
64 rho0=analytic_identity_fidelity(p,c)
65 for rho in [rho0-.02, rho0+.02]:
66 R=equicorr(p,c); T=anchored_whitener(R)
67 minf=float(metrics(R,T)[1].min())
68 transition.append({'c':c,'rho':rho,'predicted_feasible':bool(rho<=rho0),
69 'observed_feasible':bool(minf+1e-8>=rho),
70 'rho0_predicted':rho0,'min_fidelity':minf})
71 return decor, fidelity, transition
72
73
74def run_toy_classification():
75 # Same correlated standardized input for all transforms; labels depend on a rotated
76 # signal. This compares a standard per-channel normalization to exact ZCA and anchored ZCA.
77 n, p = 5000, 6
78 R = equicorr(p, .6)
79 X = rng.multivariate_normal(np.zeros(p), R, size=n)
80 w = np.array([1.0, -0.8, .5, 0, 0, 0])
81 y = (X @ w + .35*rng.normal(size=n) > 0).astype(int)
82 Xtr, Xte, ytr, yte = train_test_split(X,y,test_size=.35,random_state=SEED,stratify=y)
83 # BN-like feature standardization (using training statistics only).
84 mu=Xtr.mean(0); sd=Xtr.std(0); Xbn=(Xtr-mu)/sd; Xbnte=(Xte-mu)/sd
85 Rhat=np.cov(Xbn,rowvar=False,bias=True)
86 T=anchored_whitener(Rhat)
87 Xzca=Xbn@T; Xzcat=Xbnte@T
88 # Anchor Q=I is the identity-preserving member of the exact whitening family.
89 def fit_acc(A, At):
90 clf=LogisticRegression(C=1e3,max_iter=300,random_state=SEED)
91 clf.fit(A,ytr); return accuracy_score(yte,clf.predict(At))
92 _, fbn, offbn = metrics(Rhat,np.eye(p))
93 err_zca, fz, offz = metrics(Rhat,T)
94 return {
95 'accuracy_bn_like':fit_acc(Xbn,Xbnte),
96 'accuracy_zca_q_identity':fit_acc(Xzca,Xzcat),
97 'bn_like_offdiag_fro':float(np.linalg.norm(Rhat-np.diag(np.diag(Rhat)),'fro')),
98 'zca_offdiag_fro':float(np.linalg.norm(T.T@Rhat@T-np.diag(np.diag(T.T@Rhat@T)),'fro')),
99 'zca_max_cov_error':err_zca,
100 'zca_min_input_fidelity':float(fz.min()),
101 'rho_identity_prediction':analytic_identity_fidelity(p,.6),
102 'n_train':len(Xtr)
103 }
104
105if __name__ == '__main__':
106 decor, fidelity, transition = run_math_sweeps()
107 toy = run_toy_classification()
108 out={'seed':SEED,'predictions':{
109 'P1_exact_decorrelation':'max |T^T R T-I| should remain at floating point precision for all c and Q=I',
110 'P2_identity_fidelity':'min fidelity should equal [sqrt(1+(p-1)c)+(p-1)sqrt(1-c)]/p',
111 'P3_threshold_transition':'identity anchor is feasible iff rho <= rho0(c), where rho0 is the P2 curve'},
112 'decorrelation_sweep':decor,'fidelity_sweep':fidelity,'threshold_sweep':transition,'toy':toy}
113 print(json.dumps(out,indent=2))