From cbecb171b1af77fe6510fd60f1dfa4c7938a20d6 Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Sun, 5 Jul 2026 03:53:57 -0500 Subject: =?UTF-8?q?holofast:=20exact=20halved-jvp=20AEP=20track=20?= =?UTF-8?q?=E2=80=94=201.55x=20nudged=20phase,=20gradient-gate=20passed?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit holo_a_track computed the doubled-batch jvp/vjp on [v0; -v0] at a shared anchor zbar — exact antisymmetric redundancy (phase deviations from the common mode are exact negatives). holo_a_track_fast computes at batch B and mirrors: single-eval parity 6e-7 (exact); trajectory-level 45% divergence SHARED with the original's own FD noise floor (1e-6 state noise -> 49% self-divergence — the 2r=0.04 finite difference amplifies fp noise; training averages it via pema/momentum). Ship gate: cos(EP,BPTT) orig vs fast indistinguishable (0.907/0.912, 0.853/0.853, 0.918/0.918). Timing 6.43->4.16s on the T2=40 nudged phase (contended GPU, relative). --holofast flag, default off; queued ablation arms deliberately stay on orig for fidelity. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn --- ep_run/holo_ep.py | 44 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) (limited to 'ep_run/holo_ep.py') diff --git a/ep_run/holo_ep.py b/ep_run/holo_ep.py index 31054e4..5485bc7 100644 --- a/ep_run/holo_ep.py +++ b/ep_run/holo_ep.py @@ -254,6 +254,50 @@ def holo_a_track(blk, zs, xin, y, r, T2max, eps, K=10, exit_mult=5.0): return a_best.detach(), t_best +def holo_a_track_fast(blk, zs, xin, y, r, T2max, eps, K=10, exit_mult=5.0): + """EXACT restructure of holo_a_track (same math, ~half the correction cost): the two phase-halves' + deviations from the common mode are exact negatives (v[B:] = -v[:B], since zbar is their mean) and + the Jacobian anchor zbar is shared, so the doubled-batch jvp/vjp computed [J v0; -J v0] redundantly. + Compute Jv/JTv once at batch B and mirror. Bit-equal up to fp nondeterminism.""" + 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 = a_best = None + inc_min, t_best = 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: # measurement brake: Tikhonov-regularized adjoint + f = f - kappa * (Z - zs2a) + v0 = (Z[:B] - zbar).contiguous() # v of the +r phase; the -r phase's v is exactly -v0 + _, 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() + if inc < inc_min: + inc_min, a_best, t_best = inc, a_t, t + elif inc > exit_mult * inc_min and t >= 3 * K: + break + a_prev = a_t + if a_best is None: + a_best = a_prev if a_prev is not None else (Z[B:] - Z[:B]) / (2 * r) + t_best = T2max + return a_best.detach(), t_best + + 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 -- cgit v1.2.3