casc_eq_train v3: adaptive beta (tok-sigma stiffness compensation), adaptive-K fb with contraction test, non-contraction bail, in-training cos(EP,BP) telemetry

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn

1 files changed, 62 insertions, 15 deletions

diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py
index 99ff35c..25cfeb6 100644
--- a/ep_run/casc_eq_train.py
+++ b/ep_run/casc_eq_train.py
@@ -18,6 +18,8 @@ ap.add_argument('--K', type=int, default=3)             # fb (message-passing) r
 ap.add_argument('--geta', type=float, default=1.0)      # fb mixing (1.0 = undamped)
 ap.add_argument('--save_every', type=int, default=1000); ap.add_argument('--log', type=int, default=100)
 ap.add_argument('--wandb', default=''); ap.add_argument('--wandb_run', default='')
+ap.add_argument('--kmax', type=int, default=8)          # adaptive fb rounds cap
+ap.add_argument('--gate_every', type=int, default=200)  # in-training cos(EP,BP) telemetry
 args = ap.parse_args()
 torch.manual_seed(args.seed)
 dev = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -60,12 +62,24 @@ def free_states(x):
             z = b(z, mask); zs.append(z)
     return z0, zs
 
+@torch.no_grad()
+def tok_sigma(iters=8):
+    """top singular value of tok.weight (power iteration on the raw matrix)."""
+    W = tok.weight
+    v = torch.randn(W.shape[1], device=dev); v /= v.norm()
+    sig = 1.0
+    for _ in range(iters):
+        u = W @ v; u /= max(u.norm(), 1e-12)
+        v = W.t() @ u; sig = v.norm(); v /= max(sig, 1e-12)
+    return float(sig)
+
 def relax(z0, zs_free, y, beta):
-    """fb rounds to the nudged equilibrium: backward refresh of feedback d_l = J_{l+1}^T d_{l+1}
-    (top: -beta*NBT*dCE at current top), then forward REBUILD z_l = f_l(z_{l-1}) + d_l."""
+    """adaptive-K fb rounds: backward feedback refresh + forward rebuild, until the
+    per-round state delta contracts below 5% of round-1 (or kmax). Returns (zs, rounds, contracted)."""
     zs = [z.clone() for z in zs_free]
     d = [None] * args.L
-    for _ in range(args.K):
+    d1 = None; delta = 0.0
+    for k in range(args.kmax):
         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()
@@ -73,13 +87,22 @@ def relax(z0, zs_free, y, beta):
             zc = zs[l].detach().requires_grad_(True)
             fnext = blocks[l + 1](zc, mask)
             d[l] = torch.autograd.grad(fnext, zc, grad_outputs=d[l + 1])[0].detach()
+        delta = 0.0
         with torch.no_grad():
             prev = z0
             for l in range(args.L):
                 rebuilt = blocks[l](prev, mask) + d[l]
-                zs[l] = (1 - args.geta) * zs[l] + args.geta * rebuilt if args.geta < 1.0 else rebuilt
+                new = (1 - args.geta) * zs[l] + args.geta * rebuilt if args.geta < 1.0 else rebuilt
+                delta += float((new - zs[l]).norm())
+                zs[l] = new
                 prev = zs[l]
-    return zs
+        if k == 0:
+            d1 = max(delta, 1e-12)
+        elif delta < 0.05 * d1 and k + 1 >= args.K:
+            return zs, k + 1, True
+        elif delta > 3.0 * d1:
+            return zs, k + 1, False
+    return zs, args.kmax, delta < 0.5 * d1
 
 def dFdtheta(zs, x, y, beta):
     """dF/dtheta at fixed relaxed states (z0 rebuilt WITH graph so emb gets its E-path grad)."""
@@ -90,26 +113,38 @@ def dFdtheta(zs, x, y, beta):
     gs = torch.autograd.grad(obj, all_params, allow_unused=True)
     return [g if g is not None else None for g in gs]
 
+SIG0 = None
 def ep_step(x, y):
-    """single-sided EP: relax to the +beta equilibrium; grad = d[E/(NBT*beta) + CE]/dtheta
-    at the relaxed states (free-phase dE/dtheta == 0 exactly). Divergence guard skips bad batches."""
+    """single-sided adaptive EP: beta_t = beta0*sig0^2/sig_tok^2 (top-CE stiffness compensation),
+    adaptive-K fb relax, grad at the relaxed states. Returns (free_ce, beta_t, rounds, ok)."""
+    global SIG0
+    sig = tok_sigma()
+    if SIG0 is None: SIG0 = sig
+    beta_t = args.beta * (SIG0 * SIG0) / max(sig * sig, 1e-9)
     z0, zs_free = free_states(x)
     free_ce = F.cross_entropy(readout(zs_free[-1]).reshape(-1, vocab), y.reshape(-1)).item()
-    zp = relax(z0, zs_free, y, +args.beta)
+    zp, rounds, ok = relax(z0, zs_free, y, +beta_t)
     with torch.no_grad():
         drift = sum(float((a - b).norm()) for a, b in zip(zp, zs_free)) / max(
             sum(float(b.norm()) for b in zs_free), 1e-9)
-    if not math.isfinite(drift) or drift > 0.5:          # nudged displacement should be O(beta)
+    if (not ok) or (not math.isfinite(drift)) or drift > 0.5:
         for p in all_params: p.grad = None
-        return free_ce                                    # skip batch (guard)
+        return free_ce, beta_t, rounds, False
     prev = tok(x) + pos(torch.arange(args.T, device=dev))[None]
     E = 0.0
     for z, b in zip(zp, blocks): E = E + 0.5 * ((z - b(prev, mask)) ** 2).sum(); prev = z
-    obj = E / (NBT * args.beta) + F.cross_entropy(readout(zp[-1]).reshape(-1, vocab), y.reshape(-1))
+    obj = E / (NBT * beta_t) + F.cross_entropy(readout(zp[-1]).reshape(-1, vocab), y.reshape(-1))
     gs = torch.autograd.grad(obj, all_params, allow_unused=True)
     for p, g in zip(all_params, gs):
         p.grad = g
-    return free_ce
+    return free_ce, beta_t, rounds, True
+
+def bp_gate(x, y):
+    """true BP grads for telemetry cos (called before opt.step; reads p.grad separately)."""
+    z = tok(x) + pos(torch.arange(args.T, device=dev))[None]
+    for b in blocks: z = b(z, mask)
+    ce = F.cross_entropy(readout(z).reshape(-1, vocab), y.reshape(-1))
+    return list(torch.autograd.grad(ce, all_params, allow_unused=True))
 
 @torch.no_grad()
 def evaluate(nb=6):
@@ -134,17 +169,29 @@ n = sum(p.numel() for p in all_params)
 print(f'[{args.tag}] cascade-EP(EQUILIBRIUM/fb) L{args.L} C{args.C} T{args.T} beta={args.beta} '
       f'K={args.K} geta={args.geta} | {n/1e6:.2f}M | {dev}', flush=True)
 best, t0 = 1e9, time.time()
+skips = 0
 for step in range(args.steps + 1):
     x, y = get_batch('train')
-    ce = ep_step(x, y)
+    ce, beta_t, rounds, ok = ep_step(x, y)
+    if not ok: skips += 1
+    gcos = float('nan')
+    if step % args.gate_every == 0 and ok:
+        gbp = bp_gate(x, y)
+        num = den1 = den2 = 0.0
+        for p, g in zip(all_params, gbp):
+            if p.grad is None or g is None: continue
+            num += float((p.grad * g).sum()); den1 += float((p.grad ** 2).sum()); den2 += float((g ** 2).sum())
+        gcos = num / max((den1 ** 0.5) * (den2 ** 0.5), 1e-12)
     torch.nn.utils.clip_grad_norm_(all_params, 1.0)
     opt.step(); sched.step(); opt.zero_grad(set_to_none=True)
     if step % args.log == 0:
         val = evaluate(); best = min(best, val)
+        gtag = '' if math.isnan(gcos) else f' cos={gcos:.4f}'
         print(f'step {step:5d}/{args.steps} | train {ce:.4f} val {val:.4f} (best {best:.4f}) '
-              f'| {step/max(time.time()-t0,1e-9):.3f} it/s', flush=True)
+              f'| beta={beta_t:.2e} K={rounds} skips={skips}{gtag} | {step/max(time.time()-t0,1e-9):.3f} it/s', flush=True)
         if wb is not None:
-            try: wb.log({'train_ce': ce, 'val_ce': val, 'best': best}, step=step)
+            try: wb.log({'train_ce': ce, 'val_ce': val, 'best': best, 'beta_t': beta_t,
+                         'rounds': rounds, 'skips': skips, 'gate_cos': (None if math.isnan(gcos) else gcos)}, step=step)
             except Exception: pass
     if step % args.save_every == 0 and step > 0:
         torch.save({'tok': tok.state_dict(), 'pos': pos.state_dict(), 'blocks': blocks.state_dict(),