Sign-Reset PI Optimizer / experiment.py

Mechanism confirmed, baseline not beaten

Raw ⬇ ZIP
  1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 915
  6np.random.seed(SEED); random.seed(SEED)
  7
  8
  9def pi_run(lam, eta, kp, ki, steps=500, reset=True, x0=1.0):
 10    # Exact stated ordering: theta uses I_k; the sign event determines I_{k+1}.
 11    x = float(x0); I = 0.0; prev_g = None
 12    xs, losses, Is, resets = [], [], [], []
 13    for k in range(steps):
 14        g = lam*x
 15        event = prev_g is not None and g*prev_g < 0
 16        x = x - eta*(kp*g + ki*I)
 17        next_I = 0.0 if (reset and event) else I + g
 18        xs.append(x); losses.append(.5*lam*x*x); Is.append(next_I); resets.append(int(event and reset))
 19        I = next_I
 20        prev_g = g
 21    return np.asarray(xs), np.asarray(losses), np.asarray(Is), np.asarray(resets)
 22
 23
 24def rho(lam, eta, kp, ki):
 25    A = np.array([[1-eta*kp*lam, -eta*ki], [lam, 1.]])
 26    return max(abs(np.linalg.eigvals(A)))
 27
 28
 29def toy_checks():
 30    kp, ki, lam = 1.0, .4, 1.0
 31    predicted_eta = 4/(lam*(2*kp-ki))
 32    # Prediction 1: spectral radius reaches 1 at the Jury boundary.
 33    etas = np.linspace(.1, 3.0, 2901)
 34    rhos = np.array([rho(lam,e,kp,ki) for e in etas])
 35    crossing = etas[np.argmin(abs(rhos-1))]
 36    # Prediction 2: varying lambda scales the boundary inversely.
 37    lambdas = np.array([.5, 1., 2., 4.])
 38    boundaries = []
 39    for l in lambdas:
 40        es = np.linspace(.05, 8/l, 2401)
 41        rr = np.array([rho(l,e,kp,ki) for e in es])
 42        boundaries.append(es[np.argmin(abs(rr-1))])
 43    pred_bounds = 4/(lambdas*(2*kp-ki))
 44    # Prediction 3: reset fires on the first gradient sign reversal and zeros old memory.
 45    # Choose an intentionally oscillatory proportional step; inspect the first event.
 46    x, loss, I, reset = pi_run(1., 1.8, 1., .4, 30, True)
 47    event_idx = np.flatnonzero(reset)
 48    first = int(event_idx[0]) if len(event_idx) else -1
 49    before_I = float(I[first-1]) if first > 0 else float('nan')
 50    at_I = float(I[first]) if first >= 0 else float('nan')
 51    # Prediction 3 diagnostic: reset truncates memory at reversals, but does not
 52    # generally enlarge the linear no-reset stability region.
 53    sweep=[]
 54    for eta in np.arange(.4, 2.41, .2):
 55        row={'eta':float(eta)}
 56        for name,flag in [('reset',True),('no_reset',False)]:
 57            xx,ll,ii,rr=pi_run(1.,float(eta),kp,ki,200,flag)
 58            settle=-1
 59            for j in range(len(ll)):
 60                if np.all(ll[j:] < 1e-8): settle=j; break
 61            row[name]={'resets':int(rr.sum()),'max_abs_x':float(np.max(np.abs(xx))),
 62                       'final_loss':float(ll[-1]),'settling_step':settle}
 63        sweep.append(row)
 64    return {
 65      'predicted_eta_boundary': predicted_eta, 'observed_eta_rho1': crossing,
 66      'boundary_relative_error': abs(crossing-predicted_eta)/predicted_eta,
 67      'lambda_sweep': [{'lambda':float(l),'predicted_eta':float(p),'observed_eta':float(o),'rel_error':float(abs(o-p)/p)} for l,p,o in zip(lambdas,pred_bounds,boundaries)],
 68      'reset_first_event_step': first, 'integral_before_reset': before_I, 'integral_at_reset': at_I,
 69      'eta_sweep':sweep
 70    }
 71
 72
 73def mlp_experiment():
 74    # Small fixed synthetic two-moons-like dataset, avoiding external data downloads.
 75    try:
 76        import torch
 77        from torch import nn
 78        torch.manual_seed(SEED); np.random.seed(SEED)
 79        dev = 'cuda' if torch.cuda.is_available() else 'cpu'
 80        n=512
 81        t=np.linspace(0, math.pi, n//2)
 82        X=np.vstack([np.c_[np.cos(t),np.sin(t)], np.c_[1-np.cos(t),1-np.sin(t)-.35]])
 83        X += .08*np.random.randn(n,2)
 84        y=np.r_[np.zeros(n//2),np.ones(n//2)].astype(np.int64)
 85        perm=np.random.RandomState(SEED).permutation(n); X=X[perm]; y=y[perm]
 86        Xt=torch.tensor(X,dtype=torch.float32,device=dev); yt=torch.tensor(y,device=dev)
 87        def train(kind):
 88            torch.manual_seed(SEED+ (1 if kind=='pi' else 0))
 89            m=nn.Sequential(nn.Linear(2,24),nn.Tanh(),nn.Linear(24,2)).to(dev)
 90            lossfn=nn.CrossEntropyLoss(); prev=None; I=[torch.zeros_like(p) for p in m.parameters()]
 91            losses=[]; resets=0
 92            for step in range(300):
 93                m.zero_grad(set_to_none=True); loss=lossfn(m(Xt),yt); loss.backward()
 94                gs=[p.grad.detach().clone() for p in m.parameters()]
 95                if kind=='sgd':
 96                    with torch.no_grad():
 97                        for p,g in zip(m.parameters(),gs): p -= .08*g
 98                else:
 99                    dot=sum((g*q).sum() for g,q in zip(gs,prev)) if prev is not None else 1.
100                    if prev is not None and dot.item()<0:
101                        I=[torch.zeros_like(p) for p in m.parameters()]; resets+=1
102                    else: I=[a+g for a,g in zip(I,gs)]
103                    with torch.no_grad():
104                        for p,g,a in zip(m.parameters(),gs,I): p -= .04*(g+.4*a)
105                    prev=gs
106                losses.append(float(loss.detach().cpu()))
107            return {'final_loss':losses[-1],'final_accuracy':float((m(Xt).argmax(1)==yt).float().mean().cpu()),'resets':resets}
108        return {'device':dev,'sgd':train('sgd'),'pi_reset':train('pi')}
109    except Exception as e:
110        return {'error':repr(e),'device':'cpu'}
111
112if __name__ == '__main__':
113    out={'toy':toy_checks(),'mlp':mlp_experiment()}
114    Path('results.json').write_text(json.dumps(out,indent=2))
115    print(json.dumps(out,indent=2))