1import json, math, random, time
  2import numpy as np
  3import torch
  4
  5SEED = 1296
  6np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  7torch.set_num_threads(4)
  8device = "cuda" if torch.cuda.is_available() else "cpu"
  9try:
 10    if device == "cuda": torch.cuda.empty_cache()
 11except Exception:
 12    device = "cpu"
 13
 14def controller_step(theta, loss_fn, eta_prev, alpha=1.0, rho=.9, c=.1,
 15                    beta=.5, eta_min=1e-5, eta_max=2.0, eps=1e-12, max_bt=12):
 16    theta = theta.detach().requires_grad_(True)
 17    f0 = loss_fn(theta)
 18    g, = torch.autograd.grad(f0, theta)
 19    G = torch.linalg.vector_norm(g).item()
 20    if G < 1e-14:
 21        return theta.detach(), float(f0), 0., 0, G, 0.
 22    # Probe the currently maintained rate, then infer directional curvature.
 23    trial = theta - eta_prev * g
 24    ftrial = loss_fn(trial).detach()
 25    s_norm = eta_prev * G
 26    r = (ftrial.item() - f0.item() + eta_prev * G * G)
 27    Lhat = (1 + alpha) * max(r, 0.) / (s_norm ** (1 + alpha) + eps)
 28    eta_star = (G ** (1-alpha) / (Lhat + eps)) ** (1/alpha)
 29    eta = min(eta_max, max(eta_min, rho * eta_star))
 30    accepted = False
 31    fc = float('inf')
 32    for j in range(max_bt + 1):
 33        cand = theta - eta * g
 34        fc = loss_fn(cand).detach().item()
 35        if fc <= f0.item() - c * eta * G * G:
 36            accepted = True
 37            break
 38        eta *= beta
 39    if not accepted:
 40        eta = 0.
 41        cand = theta
 42    # Mild growth damping: preserve the accepted scale rather than blindly probing it.
 43    next_eta = min(eta_max, eta / .9) if accepted else max(eta_min, eta_prev * beta)
 44    return cand.detach(), fc if accepted else f0.item(), next_eta, int(accepted), G, Lhat
 45
 46def quadratic_sweep():
 47    rows=[]
 48    # Prediction 1: alpha=1 Lhat = directional curvature q/G^2, independent of probe eta.
 49    for lam in [0.1, 0.3, 1., 3., 10.]:
 50        H=torch.diag(torch.tensor([lam, 2.*lam, .5*lam], dtype=torch.float64))
 51        theta=torch.tensor([1.2,-.7,.4], dtype=torch.float64)
 52        def loss(x): return .5*x @ H @ x
 53        g=H@theta; G=torch.linalg.vector_norm(g).item(); q=(g @ H @ g).item()
 54        exact=q/(G*G)
 55        vals=[]
 56        for probe in [.03,.2,1.1]:
 57            trial=theta-probe*g
 58            r=(loss(trial)-loss(theta)+probe*G*G).item()
 59            lh=2*max(r,0)/(probe*G)**2
 60            vals.append(lh)
 61        rows.append({'lambda_scale':lam,'exact_directional_curvature':exact,
 62                     'observed_Lhat_mean':float(np.mean(vals)), 'relative_error':abs(np.mean(vals)-exact)/exact,
 63                     'eta_star_prediction':1/exact, 'observed_eta_star':1/float(np.mean(vals))})
 64    # Prediction 2: accepted eta * directional curvature = rho for a quadratic, and eta inversely scales.
 65    rate=[]
 66    for lam in [0.1, .3, 1., 3., 10.]:
 67        H=torch.eye(2,dtype=torch.float64)*lam; theta=torch.tensor([1.,-.5],dtype=torch.float64)
 68        def loss(x): return .5*x@H@x
 69        out=controller_step(theta,loss,eta_prev=.4,eta_max=20.,eta_min=1e-8)
 70        eta=out[2] * .9 # next_eta is eta/.9; recover accepted eta
 71        # exact expected accepted eta is rho/lam
 72        rate.append({'lambda':lam,'predicted_eta':.9/lam,'observed_eta':eta,
 73                     'eta_times_lambda':eta*lam,'accepted':bool(out[3])})
 74    # Prediction 3: sufficient-decrease acceptance boundary eta*lambda <= 2(1-c).
 75    boundary=[]; c=.1
 76    for lam in [1., 3.]:
 77        H=torch.eye(1,dtype=torch.float64)*lam; theta=torch.tensor([1.],dtype=torch.float64)
 78        def loss(x): return .5*x@H@x
 79        for eta in [1.7/lam, 1.9/lam, 2.1/lam]:
 80            g=lam*theta; f0=loss(theta).item(); fc=loss(theta-eta*g).item()
 81            accepted=fc <= f0-c*eta*(g@g).item()+1e-12
 82            boundary.append({'lambda':lam,'eta_lambda':eta*lam,'predicted_accept':eta*lam <= 2*(1-c), 'observed_accept':accepted})
 83    return {'curvature_identity':rows,'inverse_scaling':rate,'acceptance_boundary':boundary}
 84
 85def make_data(n=512):
 86    rng=np.random.RandomState(SEED)
 87    x=rng.randn(n,2).astype('float32')
 88    y=((x[:,0]*x[:,1] + .25*x[:,0] - .15*x[:,1])>0).astype('int64')
 89    return torch.tensor(x,device=device), torch.tensor(y,device=device)
 90
 91def net_loss(theta,x,y,d=16):
 92    a=2*d; b=d; c=2*d
 93    W1=theta[:a].reshape(d,2); b1=theta[a:a+b]
 94    W2=theta[a+b:a+b+c].reshape(2,d); b2=theta[a+b+c:a+b+c+2]
 95    h=torch.tanh(x@W1.T+b1)
 96    return torch.nn.functional.cross_entropy(h@W2.T+b2,y)
 97
 98def train_compare(steps=100):
 99    x,y=make_data(); d=16; n=5*d+2
100    torch.manual_seed(SEED)
101    init=torch.randn(n,device=device)*.15
102    results={}
103    for name in ['cosine_sgd','directional']:
104        theta=init.clone(); eta=.15; losses=[]; accepted=[]; lrs=[]; t0=time.perf_counter()
105        for k in range(steps):
106            fn=lambda z: net_loss(z,x,y,d)
107            if name=='directional':
108                theta,f,eta_next,ok,G,Lh=controller_step(theta,fn,eta,alpha=1.,rho=.9,c=.1,
109                                                            eta_min=1e-4,eta_max=2.,max_bt=10)
110                eta=eta_next; losses.append(f); accepted.append(ok); lrs.append(eta)
111            else:
112                th=theta.detach().requires_grad_(True); f=fn(th); g,=torch.autograd.grad(f,th)
113                lr=.15*.5*(1+math.cos(math.pi*k/steps)); theta=(th-lr*g).detach()
114                losses.append(fn(theta).item()); accepted.append(1); lrs.append(lr)
115        with torch.no_grad():
116            pred=net_loss(theta,x,y,d).item()
117            W1=theta[:2*d].reshape(d,2); b1=theta[2*d:3*d]; W2=theta[3*d:5*d].reshape(2,d); b2=theta[5*d:]
118            logits=torch.tanh(x@W1.T+b1)@W2.T+b2
119            acc=(logits.argmax(1)==y).float().mean().item()
120        results[name]={'final_loss':pred,'accuracy':acc,'loss_at_10':losses[9],
121                       'loss_at_50':losses[49],'accepted_fraction':float(np.mean(accepted)),
122                       'mean_lr':float(np.mean(lrs)),'wall_seconds':time.perf_counter()-t0}
123    return results
124
125if __name__=='__main__':
126    out={'device':device,'seed':SEED,'toy':quadratic_sweep(),'training':train_compare()}
127    with open('results.json','w') as f: json.dump(out,f,indent=2)
128    print(json.dumps(out,indent=2))