summaryrefslogtreecommitdiff
path: root/ep_run/holo_ep.py
diff options
context:
space:
mode:
authorYuren Hao <yurenh2@illinois.edu>2026-07-06 02:22:54 -0500
committerYuren Hao <yurenh2@illinois.edu>2026-07-06 02:22:54 -0500
commit8f43de671cf03ee20d5d652b6cd3ba575982a436 (patch)
tree590ab1b4ccb249d0fb15205d2480e69b7bfba080 /ep_run/holo_ep.py
parent66b4ad978585d1b8008e02d06787811e2ee93da7 (diff)
bp_lm twin control + env forensics probes + holo_a_track_avg (semi-convergence fix, ungated)
BP+EMA 2x2 (local, healthy env): qknorm {1.9951, 1.9977}, no-qknorm {1.9728, 1.9732} -> parameterization-matched BP twin tops at ~1.97 vs EP 1.7888: the equilibrium computation's iteration/depth dividend = 0.18 CE from identical parameters. External anchor (tuned depth-1 BP, 1.7921): EP at parity. bp_lm on 107/2.3.1 gave 1.9823 ~= local -> plain backward exonerated on the pascal env; the 107 divergence (pair/floss/resreg dead by step 600) narrows to the EP reg/estimator loop. Single-step fingerprints all match (a/b/c/d) -> suspected intermittent kernel issue; discriminators in flight (torch-2.7 env probe + delay_hr2's step-2000 reg-on transition). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn
Diffstat (limited to 'ep_run/holo_ep.py')
-rw-r--r--ep_run/holo_ep.py56
1 files changed, 56 insertions, 0 deletions
diff --git a/ep_run/holo_ep.py b/ep_run/holo_ep.py
index 5485bc7..d8d2be5 100644
--- a/ep_run/holo_ep.py
+++ b/ep_run/holo_ep.py
@@ -298,6 +298,62 @@ def holo_a_track_fast(blk, zs, xin, y, r, T2max, eps, K=10, exit_mult=5.0):
return a_best.detach(), t_best
+def holo_a_track_avg(blk, zs, xin, y, r, T2max, eps, K=10, exit_mult=5.0):
+ """track_fast + the semi-convergence fix (t2_probe 2026-07-05): the adjoint iteration on a
+ near-marginal operator SEMI-converges — error dips at a batch-dependent optimum then grows, and the
+ plain argmin-of-increment t_best gets fooled by rotating slow modes. Two changes:
+ (1) trend-aware stop: break after the increment rises on 2 consecutive checks past 2x inc_min
+ (instead of the blunt exit_mult=5 single-shot);
+ (2) plateau averaging: return the MEAN of the a_t whose increment <= 1.5x inc_min (the flat bottom
+ of the semi-convergence curve) — averages out the rotating error component around the optimum."""
+ import torch.func as tf
+ B = zs.size(0)
+ Z = torch.cat([zs, zs], 0)
+ X2 = torch.cat([xin, xin], 0)
+ y2 = torch.cat([y, y], 0)
+ sg = torch.cat([torch.full((B, 1, 1), r, device=zs.device), torch.full((B, 1, 1), -r, device=zs.device)], 0)
+ fnc = lambda zz: blk.nc_force(zz)
+ a_prev = None
+ hist = [] # (inc, a_t) at each K-checkpoint
+ inc_min, rise = float('inf'), 0
+ zs2a = torch.cat([zs, zs], 0)
+ kappa = getattr(blk, 'nbrake', 0.0)
+ for t in range(1, T2max + 1):
+ with torch.no_grad():
+ zbar = 0.5 * (Z[:B] + Z[B:])
+ f = rforce(blk, Z, X2) - sg * rgrad_ce(blk, Z, y2, denom=y.numel())
+ if kappa > 0:
+ f = f - kappa * (Z - zs2a)
+ v0 = (Z[:B] - zbar).contiguous()
+ _, Jv0 = tf.jvp(fnc, (zbar,), (v0,))
+ JTv0 = tf.vjp(fnc, zbar)[1](v0)[0]
+ corr0 = Jv0 - JTv0
+ Z = Z + eps * (f - torch.cat([corr0, -corr0], 0))
+ if t % K == 0 or t == T2max:
+ a_t = (Z[B:] - Z[:B]) / (2 * r)
+ if not torch.isfinite(a_t).all():
+ break
+ if a_prev is not None:
+ inc = (a_t - a_prev).norm().item()
+ hist.append((inc, a_t))
+ if inc < inc_min:
+ inc_min, rise = inc, 0
+ elif inc > 2.0 * inc_min and t >= 3 * K:
+ rise += 1 # trend-aware: need 2 consecutive rising checks
+ if rise >= 2:
+ break
+ else:
+ rise = 0
+ a_prev = a_t
+ if not hist:
+ return (a_prev if a_prev is not None else (Z[B:] - Z[:B]) / (2 * r)).detach(), T2max
+ flat = [a for inc, a in hist if inc <= 1.5 * inc_min] # the semi-convergence plateau
+ if not flat:
+ flat = [min(hist, key=lambda p: p[0])[1]]
+ a_avg = torch.stack(flat).mean(0)
+ return a_avg.detach(), len(hist) * K
+
+
def holo_a_lockin(blk, zs, xin, y, r, P, ncyc, eps):
"""True oscillatory EP / lock-in estimator (Laborieux–Zenke taken literally) — the
noisy-physics form: ONE trajectory, sinusoidal nudge beta(t)=r·sin(2πt/P), in-phase