summaryrefslogtreecommitdiff
path: root/ep_run/casc_eq_train.py
diff options
context:
space:
mode:
authorYuren Hao <yurenh2@illinois.edu>2026-07-09 15:41:54 -0500
committerYuren Hao <yurenh2@illinois.edu>2026-07-09 15:41:54 -0500
commit63cd19edce4def028f18e092d1af95ff6e0d91f4 (patch)
tree579f60cb62756e8312061e59ed2292955ce65591 /ep_run/casc_eq_train.py
parent761ac0bae654333490df7de4c6dc50e24bd863f6 (diff)
casc_eq_train: exact mode (dtop_every=1) now DEFAULT — bisection showed the v7 dedup cost ~4% CE at tuned lr; fast mode kept as opt-in. K3 verdict: exact equilibrium-EP 2.0481 vs BP 2.053+-0.004 = matched-tuning PARITY
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn
Diffstat (limited to 'ep_run/casc_eq_train.py')
-rw-r--r--ep_run/casc_eq_train.py3
1 files changed, 2 insertions, 1 deletions
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py
index 38b687e..10206f6 100644
--- a/ep_run/casc_eq_train.py
+++ b/ep_run/casc_eq_train.py
@@ -24,6 +24,7 @@ ap.add_argument('--untie', action='store_true') # separate readout matri
ap.add_argument('--tok_init', type=float, default=0.0) # >0: init tok/pos with this std (GPT-standard 0.02)
ap.add_argument('--compile', action='store_true') # torch.compile each block (free speed where supported)
ap.add_argument('--sig_every', type=int, default=25) # tok-sigma refresh interval (amortized)
+ap.add_argument('--dtop_every', type=int, default=1) # 1 = exact (DEFAULT, BP-parity); 2 = fast mode (~20% cheaper, ~4% CE tax at high lr)
ap.add_argument('--gate_every', type=int, default=200) # in-training cos(EP,BP) telemetry
args = ap.parse_args()
torch.manual_seed(args.seed)
@@ -100,7 +101,7 @@ def relax(z0, zs, ins, outs, y, beta, K, x):
reuses them instead of re-running a full graphed chain."""
d = [None] * args.L
for k in range(K):
- if k % 2 == 0 or d[args.L - 1] is None:
+ if k % args.dtop_every == 0 or d[args.L - 1] is None:
zc = zs[args.L - 1].detach().requires_grad_(True)
ce = F.cross_entropy(readout(zc).reshape(-1, vocab), y.reshape(-1))
d[args.L - 1] = (-beta * NBT * torch.autograd.grad(ce, zc)[0]).detach()