From 6ff624633ced5ff298612a8934de483a114e783e Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Thu, 9 Jul 2026 04:51:22 -0500 Subject: cascade: retire zil per user directive (EP route = equilibrium two-phase); probe gains --init_sweep (state warm-start, EP-clean readout) and --sopt adam; B2 solver sweep Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn --- ep_run/cascade_probe.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) (limited to 'ep_run') diff --git a/ep_run/cascade_probe.py b/ep_run/cascade_probe.py index da3d380..dda4059 100644 --- a/ep_run/cascade_probe.py +++ b/ep_run/cascade_probe.py @@ -24,6 +24,8 @@ ap.add_argument('--include_io', action='store_true') # gate emb/pos/tied-reado ap.add_argument('--seed', type=int, default=0); ap.add_argument('--warp', type=float, default=1.0) ap.add_argument('--ckpt', type=str, default='') # gate at a casc_bp_train checkpoint ap.add_argument('--sumscale', action='store_true') # relax in SUM units (gamma=1 Z-IL semantics) +ap.add_argument('--init_sweep', action='store_true') # one reverse gamma=1 sweep as state INIT (readout still at equilibrium = clean EP) +ap.add_argument('--sopt', choices=['sgd', 'adam'], default='sgd') # state-relaxation optimizer args = ap.parse_args() torch.manual_seed(args.seed) dev = 'cuda' if torch.cuda.is_available() else 'cpu' @@ -89,9 +91,16 @@ def relax(z0, y, beta): """minimize F_beta over states; free init = forward pass. Returns relaxed states.""" with torch.no_grad(): zs = [z.detach().clone() for z in fwd_states(z0)[1:]] + if args.init_sweep: # numerical warm-start of the STATE only + sv = args.sumscale; args.sumscale = True + for l in range(args.L - 1, -1, -1): + g = local_grad(zs, z0, l, beta, y) + zs[l] = (zs[l] - g).detach() + args.sumscale = sv if args.scheme == 'jacobi': for z in zs: z.requires_grad_(True) - opt = torch.optim.SGD(zs, lr=args.eta, momentum=args.mom) + opt = (torch.optim.Adam(zs, lr=args.eta) if args.sopt == 'adam' + else torch.optim.SGD(zs, lr=args.eta, momentum=args.mom)) for _ in range(args.K): opt.zero_grad() E = 0.0; prev = z0 @@ -180,6 +189,6 @@ print('per-block cos: ' + ' '.join(f'{l}:{np.mean(v):.4f}' for l, v in enumerat if args.include_io: idx = [i for i, n in enumerate(names) if n.startswith('io.')] print(f'io cos (last batch): {F.cosine_similarity(flat([gep[i] for i in idx]), flat([gbp[i] for i in idx]), dim=0).item():.4f}') -print(f'SUMMARY scheme={args.scheme} L={args.L} C={args.C} K={args.K} eta={args.eta} beta={args.beta} ' +print(f'SUMMARY scheme={args.scheme}{"+init" if args.init_sweep else ""}+{args.sopt} L={args.L} C={args.C} K={args.K} eta={args.eta} beta={args.beta} ' f'io={int(args.include_io)} warp={args.warp} ckpt={args.ckpt or "-"} ' f'cos={cm:.4f} cosmin={cmin:.4f} shrink={sm:.3f} t={time.time()-t0:.1f}s') -- cgit v1.2.3