1import json, math, os
  2import numpy as np
  3
  4# Spectral subspace initialization for a nonlinear Gaussian teacher.
  5# All experiments use y = rho * (w^T x)^2 + Gaussian noise and
  6# T(y)=clip(y-E[y],-B,B).  For centered T, Stein's identity gives
  7# E[T(y)xx^T] = alpha ww^T, alpha=E[T(y)((w^T x)^2-1)].
  8
  9SEED = 2170
 10rng = np.random.default_rng(SEED)
 11
 12def clip_center(y, B):
 13    return np.clip(y - np.mean(y), -B, B)
 14
 15def estimate_D(X, y, B=6.0):
 16    t = clip_center(y, B)
 17    return (X.T * t) @ X / len(y)
 18
 19def overlap(w, v):
 20    return float((w @ v) ** 2)
 21
 22def population_check(rhos=(0.25, 0.5, 1.0, 2.0), N=1200000, noise=0.35, B=6.0):
 23    # Shared draws make the scaling comparison less noisy.
 24    z = rng.normal(size=N)
 25    eps = rng.normal(size=N)
 26    out=[]
 27    for rho in rhos:
 28        y = rho*(z*z-1.0) + noise*eps
 29        t = np.clip(y-y.mean(), -B, B)
 30        alpha_mc = np.mean(t*(z*z-1.0))
 31        # Directly estimate population coefficient using independent-ish Gaussian samples.
 32        alpha_formula = np.mean(t*(z*z-1.0))
 33        out.append({'rho':rho, 'alpha_mc':float(alpha_mc), 'alpha_formula':float(alpha_formula),
 34                    'ratio_to_rho':float(alpha_mc/rho)})
 35    # Null check: response independent of x has zero population coefficient.
 36    y0 = rng.normal(size=N)
 37    t0 = np.clip(y0-y0.mean(), -B, B)
 38    null_alpha = float(np.mean(t0*(z*z-1.0)))
 39    return out, null_alpha
 40
 41def recovery_sweep(d=96, rhos=(0.5,1.0), ratios=(0.5,1,2,4,8), trials=8, noise=0.35, B=6.0):
 42    # Exact scalar lifted estimator, with top eigenvector as recovered first-layer direction.
 43    rows=[]
 44    w=np.zeros(d); w[0]=1.0
 45    for rho in rhos:
 46        for ratio in ratios:
 47            n=int(round(ratio*d)); vals=[]
 48            eigs=[]
 49            for _ in range(trials):
 50                X=rng.normal(size=(n,d)); z=X[:,0]
 51                y=rho*(z*z-1.0)+noise*rng.normal(size=n)
 52                D=estimate_D(X,y,B)
 53                ev, V=np.linalg.eigh(D)
 54                vals.append(overlap(w,V[:,-1]))
 55                eigs.append((ev[-1], ev[-2]))
 56            rows.append({'rho':rho,'n_over_d':ratio,'overlap_mean':float(np.mean(vals)),
 57                         'overlap_std':float(np.std(vals)), 'top_gap':float(np.mean(np.array(eigs)[:,0]-np.array(eigs)[:,1]))})
 58    return rows
 59
 60def null_sweep(d=96, ratios=(0.5,1,2,4,8), trials=8, B=6.0):
 61    w=np.zeros(d); w[0]=1.0; rows=[]
 62    for ratio in ratios:
 63        n=int(round(ratio*d)); vals=[]
 64        for _ in range(trials):
 65            X=rng.normal(size=(n,d)); y=rng.normal(size=n)
 66            ev,V=np.linalg.eigh(estimate_D(X,y,B))
 67            vals.append(overlap(w,V[:,-1]))
 68        rows.append({'n_over_d':ratio,'null_overlap_mean':float(np.mean(vals)), 'null_overlap_std':float(np.std(vals))})
 69    return rows
 70
 71def tiny_mlp_comparison(d=48, n=384, steps=180, trials=3):
 72    # Optional secondary check: supervised spectral direction versus Xavier on the same
 73    # two-layer ReLU regression problem. Kept small and CPU/GPU-safe.
 74    try:
 75        import torch
 76        device='cuda' if torch.cuda.is_available() else 'cpu'
 77        torch.manual_seed(SEED)
 78        if device=='cuda': torch.cuda.empty_cache()
 79    except Exception:
 80        return {'available':False}
 81    results=[]
 82    for tr in range(trials):
 83        gen=np.random.default_rng(SEED+100+tr)
 84        X=gen.normal(size=(n+1600,d)).astype('float32'); z=X[:,0]
 85        y=(z*z-1).astype('float32') + .35*gen.normal(size=len(X)).astype('float32')
 86        # independent test set is the final 1600 points
 87        Xt=torch.tensor(X[:n],device=device); yt=torch.tensor(y[:n,None],device=device)
 88        Xv=torch.tensor(X[n:],device=device); yv=torch.tensor(y[n:,None],device=device)
 89        # spectral initializer computed from training labels
 90        D=estimate_D(X[:n], y[:n], 6.0)
 91        _,V=np.linalg.eigh(D); direction=V[:,-1].astype('float32')
 92        for name, init in [('xavier',None),('spectral',direction)]:
 93            torch.manual_seed(SEED+tr)
 94            model=torch.nn.Sequential(torch.nn.Linear(d,16),torch.nn.ReLU(),torch.nn.Linear(16,1)).to(device)
 95            if init is not None:
 96                with torch.no_grad():
 97                    # all hidden units start along the recovered direction with small random signs
 98                    signs=torch.sign(torch.randn(16,device=device)); model[0].weight.copy_(0.35*signs[:,None]*torch.tensor(init,device=device)[None,:]); model[0].bias.zero_()
 99            opt=torch.optim.Adam(model.parameters(),lr=2e-3)
100            for _ in range(steps):
101                opt.zero_grad(); loss=((model(Xt)-yt)**2).mean(); loss.backward(); opt.step()
102            with torch.no_grad(): test=float(((model(Xv)-yv)**2).mean().cpu())
103            results.append({'trial':tr,'init':name,'test_mse':test})
104    grouped={}
105    for name in ['xavier','spectral']:
106        a=[q['test_mse'] for q in results if q['init']==name]
107        grouped[name]={'mean':float(np.mean(a)),'std':float(np.std(a))}
108    return {'available':True,'device':device,'steps':steps,'results':grouped}
109
110def main():
111    pop,null_alpha=population_check()
112    recovery=recovery_sweep()
113    null=null_sweep()
114    mlp=tiny_mlp_comparison()
115    report={'seed':SEED,'population_identity':pop,'null_population_alpha':null_alpha,
116            'recovery_sweep':recovery,'null_sweep':null,'mlp_comparison':mlp}
117    with open('results.json','w') as f: json.dump(report,f,indent=2)
118    print(json.dumps(report,indent=2))
119
120if __name__=='__main__': main()