summaryrefslogtreecommitdiff
path: root/ep_run/spec_rho_vs_c.py
diff options
context:
space:
mode:
authorYuren Hao <yurenh2@illinois.edu>2026-07-03 05:56:50 -0500
committerYuren Hao <yurenh2@illinois.edu>2026-07-03 05:56:50 -0500
commitb83947778e2c776f757a07d4719b7ce961d7ed55 (patch)
treeb9cc01d7adda691d9156d9d04f4fb2f644674e96 /ep_run/spec_rho_vs_c.py
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
Diffstat (limited to 'ep_run/spec_rho_vs_c.py')
-rw-r--r--ep_run/spec_rho_vs_c.py29
1 files changed, 29 insertions, 0 deletions
diff --git a/ep_run/spec_rho_vs_c.py b/ep_run/spec_rho_vs_c.py
new file mode 100644
index 0000000..47661e0
--- /dev/null
+++ b/ep_run/spec_rho_vs_c.py
@@ -0,0 +1,29 @@
+import torch, math, pickle
+from pathlib import Path
+import lt_ep_train as L
+from lt_ep_train import EQBlock
+L.DD=Path('data/tinystories_bpe'); L.vocab=pickle.load(open(L.DD/'meta.pkl','rb'))['vocab_size']
+dev='cuda'; eps=0.1; B=8; T=256; N=800
+torch.manual_seed(1234); idx,y=L.get_batch('val',B,T)
+idx=idx.to(dev) if hasattr(idx,'to') else idx
+def measure_c(ckpt,c):
+ blk=EQBlock(512,16,256,256,s=1.0,c=c,attn_mode='thick'); blk.qknorm=True
+ ck=torch.load(ckpt,map_location=dev)
+ with torch.no_grad():
+ for p,w in zip(blk.allp,ck['allp']): p.copy_(w.to(dev))
+ xin=blk.embed(idx).detach(); z=xin.clone(); ress=[]
+ for t in range(N):
+ z2=z+eps*blk.force(z,xin).detach()
+ r=(z2-z).norm().item()/(z.norm().item()+1e-9); ress.append(r); z=z2
+ if not math.isfinite(r) or r>1e2: break
+ win=[ress[i] for i in range(len(ress)) if 1e-6<ress[i]<1e-1] or ress[-200:]
+ rats=[win[i+1]/win[i] for i in range(len(win)-1) if win[i]>0]
+ rho=math.exp(sum(math.log(x) for x in rats)/len(rats)) if rats else float('nan')
+ return rho, ress[-1]
+print("=== rho vs damping c — does more c pull the operator off the rho=1 threshold? ===")
+print("(weights trained at c=1; this is eval-time c — a margin indicator, not the trained answer)")
+for ck,lab in [('runs/redx_traj/s3200.pt','redx-s3200 (val2.74, marginal)'),('runs/bptt_final.pt','BPTT (1.83)')]:
+ for c in [1.0,1.5,2.0,3.0,4.0]:
+ try: rho,fr=measure_c(ck,c); print(f" {lab} c={c}: rho={rho:.4f} final_res={fr:.2e}")
+ except Exception as e: print(f" {lab} c={c}: ERR {repr(e)[:60]}")
+print("=== DONE ===")