Anchored Whitening Layer / experiment.py

Unverified

Raw ⬇ ZIP
  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))