Transfer-Spectrum Pseudo-Transition Scheduler / experiment.py

Mechanism confirmed, baseline not beaten

Raw ⬇ ZIP
  1import os, json, math, random
  2import numpy as np
  3
  4# Reproducible toy verification plus a small optimizer experiment.
  5SEED = 7
  6np.random.seed(SEED); random.seed(SEED)
  7
  8
  9def avoided(a, b, g):
 10    m = (a+b)/2
 11    d = np.sqrt(((a-b)/2)**2 + g*g)
 12    return m+d, m-d
 13
 14
 15def toy_check():
 16    # Two uncoupled modes cross at theta=0: a,b have slopes +/-s.
 17    s, base = 0.8, 1.0
 18    thetas = np.linspace(-2, 2, 4001)
 19    rows=[]
 20    for g in [0.02, 0.05, 0.10, 0.20]:
 21        a = base+s*thetas; b = base-s*thetas
 22        lp,lm = avoided(a,b,g)
 23        # Positive eigenvalues in the central region; normalized gap as proposed.
 24        G=(lp-np.abs(lm))/lp
 25        i=int(np.argmin(G))
 26        # xi based on the second eigenvalue, and response d(lambda+)/d(theta).
 27        xi=1/np.log(lp/np.maximum(np.abs(lm),1e-12))
 28        # Exact curvature of lambda_plus: s^2*g^2 / ((s*theta)^2+g^2)^(3/2), maximal at resonance.
 29        response=(s*s*g*g)/((s*thetas)**2+g*g)**1.5
 30        ir=int(np.argmax(response))
 31        # Predictions: crossing/response at 0, splitting 2g, xi(0) ~ base/(2g).
 32        split=lp[i]-lm[i]
 33        rows.append(dict(g=g, theta_gap=float(thetas[i]), predicted_theta=0.,
 34                         split=float(split), predicted_split=2*g,
 35                         xi_at_cross=float(xi[np.argmin(np.abs(thetas))]),
 36                         predicted_xi=base/(2*g),
 37                         theta_response=float(thetas[ir])))
 38    # scaling fit xi vs 1/g, and normalized errors
 39    gs=np.array([r['g'] for r in rows]); xis=np.array([r['xi_at_cross'] for r in rows])
 40    slope=float(np.polyfit(1/gs,xis,1)[0])
 41    # Correlation prediction: a pure subleading mode has C(k)/C(0)=(|lambda1|/lambda0)^k.
 42    g=.1; a=base+s*thetas; b=base-s*thetas; lp,lm=avoided(a,b,g)
 43    j=np.argmin(np.abs(thetas)); ratio=abs(lm[j])/lp[j]
 44    ks=np.arange(1,11); corr=ratio**ks
 45    fit=float(np.polyfit(ks,np.log(corr),1)[0])
 46    memory_xi=float(-1/fit)
 47    return rows, slope, memory_xi
 48
 49
 50def make_data(n, seq=20):
 51    x=np.random.rand(n,seq,2).astype('float32')
 52    y=x.sum((1,2)).astype('float32')
 53    return x,y
 54
 55# Torch is imported lazily so the mathematical check also works without CUDA.
 56def train(use_scheduler, seed=7, steps=500):
 57    import torch
 58    torch.manual_seed(seed); np.random.seed(seed)
 59    device='cuda' if torch.cuda.is_available() else 'cpu'
 60    try:
 61        class RNN(torch.nn.Module):
 62            def __init__(self):
 63                super().__init__(); self.cell=torch.nn.RNNCell(2,24,nonlinearity='tanh'); self.out=torch.nn.Linear(24,1)
 64            def forward(self,x, collect=False):
 65                h=torch.zeros(x.shape[0],24,device=x.device); hs=[]
 66                for t in range(x.shape[1]):
 67                    h=self.cell(x[:,t],h); hs.append(h)
 68                return self.out(h).squeeze(1), hs
 69        model=RNN().to(device)
 70        opt=torch.optim.Adam(model.parameters(),lr=0.01)
 71        lr=0.01; records=[]
 72        for step in range(steps):
 73            x,y=make_data(64,20); x=torch.tensor(x,device=device); y=torch.tensor(y,device=device)
 74            pred,hs=model(x,True); loss=((pred-y)**2).mean()
 75            opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5.0); opt.step()
 76            gap=entropy=xi=np.nan
 77            if use_scheduler and step%10==0 and step>=10:
 78                # Fit h_{t+1}=T h_t over batch and time, as the proposal suggests.
 79                A=torch.stack(hs[:-1],1).detach().reshape(-1,24)
 80                B=torch.stack(hs[1:],1).detach().reshape(-1,24)
 81                # T maps row states: A @ M ~= B; eigenvalues of M.
 82                M=torch.linalg.lstsq(A,B,rcond=None).solution
 83                ev=torch.linalg.eigvals(M).real.abs().sort(descending=True).values
 84                l0=max(float(ev[0]),1e-8); l1=float(ev[1]) if len(ev)>1 else 0
 85                gap=(l0-l1)/l0; xi=1/max(math.log(l0/max(l1,1e-8)),1e-8)
 86                # Projection entropy onto leading right eigenvectors, normalized.
 87                vals,vecs=torch.linalg.eig(M); ix=torch.argsort(vals.real.abs(),descending=True)[:3]
 88                V=vecs[:,ix].real; z=A@V; p=(z*z).mean(0); p=p/(p.sum()+1e-8)
 89                entropy=float(-(p*torch.log(p+1e-8)).sum())
 90                if gap<0.12 and entropy>0.7:
 91                    lr=max(lr*0.5,0.0005)
 92                else:
 93                    lr=float(np.clip(lr*(gap/0.25),0.0005,0.01))
 94                for pg in opt.param_groups: pg['lr']=lr
 95            records.append((float(loss.detach().cpu()), gap, entropy, lr, xi))
 96        tail=np.array([r[0] for r in records[-50:]])
 97        return dict(device=device, final=float(records[-1][0]), best=float(np.min([r[0] for r in records])), tail=float(tail.mean()),
 98                    lr_final=float(lr), detections=int(sum(np.isfinite(r[1]) and r[1]<.12 and r[2]>.7 for r in records)), records=records)
 99    except Exception as e:
100        if device=='cuda':
101            # Retry the identical experiment on CPU if CUDA is unavailable/fragile.
102            torch.cuda.empty_cache()
103            original=torch.cuda.is_available
104            torch.cuda.is_available=lambda: False
105            try:
106                return train(use_scheduler, seed, steps)
107            finally:
108                torch.cuda.is_available=lambda: original
109        raise
110
111if __name__=='__main__':
112    rows,slope,memory_xi=toy_check()
113    out={'toy':rows,'xi_slope_vs_1_over_g':slope,'memory_xi_numeric':memory_xi,'memory_xi_predicted':rows[2]['xi_at_cross']}
114    try:
115        out['baseline']=train(False)
116        out['idea']=train(True)
117    except Exception as e:
118        out['training_error']=repr(e)
119    with open('results.json','w') as f: json.dump(out,f,indent=2)
120    print(json.dumps(out,indent=2))