holofast: exact halved-jvp AEP track — 1.55x nudged phase, gradient-gate passed

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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn

1 files changed, 44 insertions, 0 deletions

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