1import json, math, random
  2import numpy as np
  3import torch
  4
  5
  6def ridge(S, b, lam):
  7    return np.linalg.solve(S + lam*np.eye(S.shape[0]), b)
  8
  9
 10def estimate_lambda(w, S, b, lo=1e-8, hi=10.0):
 11    grid = np.r_[0.0, np.logspace(math.log10(lo), math.log10(hi), 240)]
 12    scores = np.array([np.sum((w-ridge(S,b,lam))**2) for lam in grid])
 13    j = int(np.argmin(scores))
 14    if 0 < j < len(grid)-1:
 15        a, z = max(grid[j-1], grid[1]), grid[j+1]
 16        for lam in np.logspace(np.log10(a), np.log10(z), 80):
 17            score=np.sum((w-ridge(S,b,lam))**2)
 18            if score < scores[j]: scores[j], grid[j] = score, lam
 19    return float(grid[j])
 20
 21
 22def math_check(seed=7):
 23    rng=np.random.default_rng(seed); d=8; n=30000
 24    A=rng.normal(size=(d,d)); S=A@A.T/d+.2*np.eye(d)
 25    theta=rng.normal(size=d); true_lam=.37
 26    X=rng.multivariate_normal(np.zeros(d),S,n); y=X@theta
 27    b=S@theta; w=ridge(S,b,true_lam)
 28    Sh=np.zeros((d,d)); bh=np.zeros(d); alpha=.97
 29    for i in range(0,n,100):
 30        H=X[i:i+100]; yy=y[i:i+100]
 31        Sh=alpha*Sh+(1-alpha)*H.T@H/len(H)
 32        bh=alpha*bh+(1-alpha)*H.T@yy/len(H)
 33    est=estimate_lambda(w,Sh,bh,hi=4)
 34    return {'true_lambda':true_lam,'estimated_lambda':est,
 35            'abs_error':abs(est-true_lam),
 36            'ridge_residual':float(np.linalg.norm(w-ridge(S,b,true_lam)))}
 37
 38
 39def run(seed, mode, steps=500):
 40    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
 41    device='cuda' if torch.cuda.is_available() else 'cpu'
 42    try:
 43        d,N,batch=12,1800,64; eps=.18; lr=.035
 44        gen=np.random.default_rng(seed); scales=np.linspace(.45,2.0,d)
 45        X=gen.normal(size=(N,d))*scales; theta=gen.normal(size=d); theta/=np.linalg.norm(theta)
 46        y=X@theta+.15*gen.normal(size=N)
 47        Xt=torch.tensor(X,dtype=torch.float32,device=device); yt=torch.tensor(y,dtype=torch.float32,device=device)
 48        w=torch.zeros(d,device=device,requires_grad=True)
 49        wd=.035 if mode=='fixed_weight_decay' else 0.
 50        opt=torch.optim.SGD([w],lr=lr,weight_decay=wd)
 51        Sema=torch.zeros((d,d),device=device); bema=torch.zeros(d,device=device)
 52        alpha=.96; lambda0=.035; rho=.28; smoothed=0.; updates=0; hist=[]
 53        for t in range(steps):
 54            ix=torch.randint(0,N,(batch,),device=device); H=Xt[ix]; yy=yt[ix]
 55            opt.zero_grad(set_to_none=True); residual=H@w-yy
 56            loss=.5*(torch.abs(residual)+eps*torch.linalg.vector_norm(w)).square().mean()
 57            loss.backward()
 58            with torch.no_grad():
 59                Sema=alpha*Sema+(1-alpha)*H.T@H/batch
 60                bema=alpha*bema+(1-alpha)*H.T@yy/batch
 61                if mode=='controller' and t>=60 and t%10==0:
 62                    Sn=Sema.cpu().numpy(); bn=bema.cpu().numpy(); wn=w.cpu().numpy()
 63                    current=estimate_lambda(wn,Sn,bn,hi=3.)
 64                    smoothed=current if updates==0 else .85*smoothed+.15*current
 65                    # Do not react to a >10x jump in the noisy early estimate.
 66                    if updates==0 or (.1*max(smoothed,1e-8)<=current<=10*max(smoothed,1e-8)):
 67                        updates+=1; hist.append(smoothed)
 68                if mode=='controller' and updates>=2:
 69                    w.grad.add_(rho*(smoothed-lambda0)*w)
 70            opt.step()
 71        with torch.no_grad():
 72            r=Xt@w-yt; clean=.5*r.square().mean(); robust=.5*(torch.abs(r)+eps*torch.linalg.vector_norm(w)).square().mean()
 73            return {'clean_loss':float(clean.cpu()),'robust_loss':float(robust.cpu()),
 74                    'weight_norm':float(torch.linalg.vector_norm(w).cpu()),
 75                    'lambda_last':float(hist[-1]) if hist else None,'device':device}
 76    except Exception:
 77        if device=='cuda':
 78            torch.cuda.empty_cache()
 79            # Explicit CPU retry, preserving deterministic seed and settings.
 80            return run_cpu(seed,mode,steps)
 81        raise
 82
 83
 84def run_cpu(seed,mode,steps):
 85    # CUDA fallback uses the same numerical path on CPU.
 86    available=torch.cuda.is_available
 87    torch.cuda.is_available=lambda:False
 88    try: return run(seed,mode,steps)
 89    finally: torch.cuda.is_available=available
 90
 91
 92def main():
 93    out={'math_check':math_check(),'runs':{}}
 94    for mode in ['adv_sgd','fixed_weight_decay','controller']:
 95        vals=[run(s,mode) for s in [11,23,41,59,71]]; out['runs'][mode]=vals
 96        out['summary_'+mode]={
 97            'clean_loss':float(np.mean([v['clean_loss'] for v in vals])),
 98            'robust_loss':float(np.mean([v['robust_loss'] for v in vals])),
 99            'weight_norm':float(np.mean([v['weight_norm'] for v in vals])),
100            'robust_std':float(np.std([v['robust_loss'] for v in vals]))}
101    print(json.dumps(out,indent=2))
102
103if __name__=='__main__': main()