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, 49 insertions, 0 deletions

diff --git a/ep_run/resreg_warm_probe_loop.py b/ep_run/resreg_warm_probe_loop.py
new file mode 100644
index 0000000..21c73a2
--- /dev/null
+++ b/ep_run/resreg_warm_probe_loop.py
@@ -0,0 +1,49 @@
+import time, os, re, torch, pickle, numpy as np, subprocess
+from pathlib import Path
+from scipy.sparse.linalg import eigs, LinearOperator
+import lt_ep_train as L
+from lt_ep_train import EQBlock, relax
+os.chdir("/home/yurenh2/ept/ep_run")
+L.DD=Path('data/tinystories_bpe'); L.vocab=pickle.load(open(L.DD/'meta.pkl','rb'))['vocab_size']
+dev='cuda'; B=24; T=256; eps=0.1
+CK='runs/ep_resreg_warm.pt'; LOG='runs/ep_resreg_warm.log'; OUT='runs/resreg_warm_probe.log'
+torch.manual_seed(0); idx,y=L.get_batch('val',B,T); idx=idx.to(dev) if hasattr(idx,'to') else idx
+
+def alive(): return subprocess.run(["pgrep","-f","ckpt runs/ep_resreg_warm.pt"],capture_output=True).returncode==0
+def curinfo():
+    try:
+        ls=[l for l in open(LOG) if l.startswith("step")]
+        m=re.search(r"step\s+(\d+)",ls[-1]); b=re.search(r"best ([\d.]+)",ls[-1]); v=re.search(r"val CE ([\d.]+)",ls[-1])
+        return (m.group(1), b.group(1), v.group(1))
+    except Exception: return ("?","?","?")
+def probe():
+    ck=torch.load(CK,map_location=dev)
+    blk=EQBlock(512,16,256,256,s=1.0,c=1.0,attn_mode='thick'); blk.qknorm=True; blk.fnoise=0.0
+    with torch.no_grad():
+        for p,w in zip(blk.allp,ck['allp']): p.copy_(w.to(dev))
+        xin=blk.embed(idx).detach()
+        z=relax(blk,xin.clone(),xin,300,eps).detach()
+        g=(blk.force(z,xin).detach().norm()/z.norm()).item()
+        F0=blk.force(z,xin).detach(); h=1e-3*z.norm().item(); N=z.numel()
+        def Jv(v):
+            vt=torch.from_numpy(v).float().to(dev).view_as(z); nv=vt.norm().item()
+            if nv<1e-20: return np.zeros_like(v)
+            return ((blk.force(z+h*(vt/nv),xin).detach()-F0)/h*nv).view(-1).cpu().numpy()
+        lam=eigs(LinearOperator((N,N),matvec=lambda v:v+eps*Jv(v),dtype=np.float32),k=6,which='LM',return_eigenvectors=False,maxiter=300,tol=1e-3)
+        mu=(lam-1)/eps; top=mu[np.argmax(mu.real)]
+        return g, float(top.real), float(abs(top.imag))
+
+t0=time.time()
+hdr="=== resreg_warm g/Reμ probe loop (every 20min) — track if g holds (bounded) or grows (toward blow) ==="
+print(hdr,flush=True); open(OUT,'a').write(hdr+"\n")
+while time.time()-t0 < 24*3600:
+    if not alive():
+        line=f"[+{int(time.time()-t0)}s] ep_resreg_warm DEAD/done -> stop"; print(line,flush=True); open(OUT,'a').write(line+"\n"); break
+    try:
+        st,bb,vv=curinfo(); g,re_mu,im_mu=probe()
+        line=f"[+{int((time.time()-t0)/60)}min] step{st} best{bb} valCE{vv} | g_floor={g:.2e} Re_mu={re_mu:+.3f} Im={im_mu:.3f}"
+        print(line,flush=True); open(OUT,'a').write(line+"\n")
+    except Exception as e:
+        print("probe skip:",repr(e),flush=True)
+    time.sleep(1200)
+print("loop end",flush=True)