Initial commit: ept — backprop-free equilibrium transformer (EP)

Code (ep_run/), organized docs (docs/{method,campaign,hardware,outreach,paper}),
analysis scripts (scripts/), ONBOARDING.md entry point. Large data/checkpoints
git-ignored (share separately).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn

1 files changed, 43 insertions, 0 deletions

diff --git a/ep_run/eig_probe.py b/ep_run/eig_probe.py
new file mode 100644
index 0000000..a87864b
--- /dev/null
+++ b/ep_run/eig_probe.py
@@ -0,0 +1,43 @@
+import torch, pickle, numpy as np
+from pathlib import Path
+from scipy.sparse.linalg import LinearOperator, eigs
+import lt_ep_train as L
+from lt_ep_train import EQBlock, relax
+L.DD=Path('data/tinystories_bpe'); L.vocab=pickle.load(open(L.DD/'meta.pkl','rb'))['vocab_size']
+dev='cuda'; B=2; T=256; eps=0.1
+torch.manual_seed(1234); idx,y=L.get_batch('val',B,T); idx=idx.to(dev) if hasattr(idx,'to') else idx
+def load_blk(path):
+    ck=torch.load(path,map_location=dev)
+    blk=EQBlock(512,16,256,256,s=1.0,c=1.0,attn_mode='thick'); blk.qknorm=True
+    with torch.no_grad():
+        for p,w in zip(blk.allp,ck['allp']): p.copy_(w.to(dev))
+    return blk, ck.get('best','?'), ck.get('step','?')
+@torch.no_grad()
+def lead_eigs(blk, relax_steps=250, k=8, hrel=1e-3):
+    xin=blk.embed(idx).detach(); z=relax(blk,xin.clone(),xin,relax_steps,eps).detach()
+    F0=blk.force(z,xin).detach(); zn=z.norm().item(); gfloor=F0.norm().item()/(zn+1e-9)
+    shp=z.shape; N=z.numel(); h=hrel*zn
+    def Jv(vt):                      # finite-difference JVP: J@vt
+        nv=vt.norm().item()
+        if nv<1e-20: return torch.zeros_like(vt)
+        pert=h*(vt/nv)
+        return (blk.force(z+pert,xin).detach()-F0)/h*nv
+    # sanity: random-direction ||Jv||
+    vr=torch.randn(shp,device=dev); jr=Jv(vr); sane=jr.norm().item()/(vr.norm().item()+1e-9)
+    def matvec(v):
+        vt=torch.from_numpy(np.ascontiguousarray(v).astype('float32')).reshape(shp).to(dev)
+        return (vt+eps*Jv(vt)).double().cpu().numpy().reshape(-1)
+    op=LinearOperator((N,N),matvec=matvec,dtype='float64')
+    vals=eigs(op,k=k,which='LM',return_eigenvectors=False,maxiter=4000,tol=1e-5)
+    return gfloor, sane, sorted(vals,key=lambda x:-abs(x))
+for tag,path in [("s2000 (healthy)","runs/redx_traj/s2000.pt"),
+                 ("s3200 (blew@2.74)","runs/redx_traj/s3200.pt"),
+                 ("ep_eps05 (blew@2.41)","runs/ep_eps05.pt")]:
+    blk,best,step=load_blk(path)
+    g,sane,vals=lead_eigs(blk)
+    print(f"=== {tag} best={best} step={step} g_floor={g:.4f} ||Jv_rand||/||v||={sane:.3f} ===")
+    for lam in vals[:5]:
+        mu=(lam-1)/eps
+        fl=("UNSTABLE" if abs(lam)>1.0+1e-4 else "stable")+(" CMPLX" if abs(lam.imag)>1e-3 else " real")+(" ReMu<0=Euler-artifact" if mu.real<-1e-4 else " ReMu>=0=TRUE-instab")
+        print(f"   |lam|={abs(lam):.5f} lam={lam.real:+.4f}{lam.imag:+.4f}j mu={mu.real:+.4f}{mu.imag:+.4f}j [{fl}]")
+print("=== DONE ===")