cascade root cause found: tied readout + default N(0,1) embedding init = pathological top-CE stiffness once predictions sharpen (sigma_tok~76 vs GPT-standard 0.02 giving ~1.6); add --untie and --tok_init; diag arms A/B/C

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

2 files changed, 23 insertions, 8 deletions

diff --git a/ep_run/casc_bp_train.py b/ep_run/casc_bp_train.py
index 8e039ec..5194c45 100644
--- a/ep_run/casc_bp_train.py
+++ b/ep_run/casc_bp_train.py
@@ -14,6 +14,7 @@ ap.add_argument('--lr', type=float, default=3e-4); ap.add_argument('--warmup', t
 ap.add_argument('--seed', type=int, default=0)
 ap.add_argument('--save_every', type=int, default=500); ap.add_argument('--log', type=int, default=200)
 ap.add_argument('--wandb', default=''); ap.add_argument('--wandb_run', default='')
+ap.add_argument('--tok_init', type=float, default=0.0)  # >0: init tok/pos std (GPT-standard 0.02)
 args = ap.parse_args()
 torch.manual_seed(args.seed)
 dev = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -40,6 +41,9 @@ class Block(nn.Module):
 
 tok = nn.Embedding(vocab, args.C).to(dev)
 pos = nn.Embedding(args.T, args.C).to(dev)
+if args.tok_init > 0:
+    with torch.no_grad():
+        tok.weight.normal_(0, args.tok_init); pos.weight.normal_(0, args.tok_init)
 blocks = nn.ModuleList([Block(args.C, args.H) for _ in range(args.L)]).to(dev)
 mask = torch.triu(torch.full((args.T, args.T), float('-inf'), device=dev), 1)
 params = list(tok.parameters()) + list(pos.parameters()) + list(blocks.parameters())
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py
index 96b7ab3..547312e 100644
--- a/ep_run/casc_eq_train.py
+++ b/ep_run/casc_eq_train.py
@@ -19,6 +19,9 @@ ap.add_argument('--geta', type=float, default=1.0)      # fb mixing (1.0 = undam
 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('--noguard', action='store_true')       # diagnosis: skip only non-finite grads
+ap.add_argument('--untie', action='store_true')         # separate readout matrix (untied from tok)
+ap.add_argument('--tok_init', type=float, default=0.0)  # >0: init tok/pos with this std (GPT-standard 0.02)
 ap.add_argument('--gate_every', type=int, default=200)  # in-training cos(EP,BP) telemetry
 args = ap.parse_args()
 torch.manual_seed(args.seed)
@@ -46,10 +49,14 @@ class Block(nn.Module):
 
 tok = nn.Embedding(vocab, args.C).to(dev)
 pos = nn.Embedding(args.T, args.C).to(dev)
+if args.tok_init > 0:
+    with torch.no_grad():
+        tok.weight.normal_(0, args.tok_init); pos.weight.normal_(0, args.tok_init)
 blocks = nn.ModuleList([Block(args.C, args.H) for _ in range(args.L)]).to(dev)
 mask = torch.triu(torch.full((args.T, args.T), float('-inf'), device=dev), 1)
-readout = lambda z: z @ tok.weight.t()
-all_params = list(tok.parameters()) + list(pos.parameters()) + list(blocks.parameters())
+W_out = nn.Parameter(torch.randn(vocab, args.C, device=dev) * 0.02) if args.untie else None
+readout = (lambda z: z @ W_out.t()) if args.untie else (lambda z: z @ tok.weight.t())
+all_params = list(tok.parameters()) + list(pos.parameters()) + list(blocks.parameters()) + ([W_out] if args.untie else [])
 opt = torch.optim.AdamW(all_params, lr=args.lr, weight_decay=1e-4)
 sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1.0, (s + 1) / max(args.warmup, 1)))
 NBT = args.B * args.T
@@ -65,7 +72,7 @@ def free_states(x):
 @torch.no_grad()
 def tok_sigma(iters=8):
     """top singular value of tok.weight (power iteration on the raw matrix)."""
-    W = tok.weight
+    W = W_out if args.untie else tok.weight
     v = torch.randn(W.shape[1], device=dev); v /= v.norm()
     sig = 1.0
     for _ in range(iters):
@@ -104,13 +111,14 @@ def dFdtheta(zs, x, y, beta):
     return [g if g is not None else None for g in gs]
 
 SIG0 = None
-GOV = {'K': None, 'bscale': 1.0, 'gema': None}
+GOV = {'K': None, 'bscale': 1.0, 'gema': None, 'drift': 0.0, 'gn': 0.0, 'sig': 0.0}
 def ep_step(x, y):
     """single-sided EP with a QUALITY-GOVERNED estimator: beta_t = beta0*bscale*sig0^2/sig^2,
     K = GOV['K'] fb rounds; guard = finiteness + drift + grad-norm sanity only."""
     global SIG0
     if GOV['K'] is None: GOV['K'] = args.K
     sig = tok_sigma()
+    GOV['sig'] = sig
     if SIG0 is None: SIG0 = sig
     beta_t = args.beta * GOV['bscale'] * (SIG0 * SIG0) / max(sig * sig, 1e-9)
     z0, zs_free = free_states(x)
@@ -119,9 +127,10 @@ def ep_step(x, y):
     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:
+    if (not math.isfinite(drift)) or (drift > 0.5 and not args.noguard):
         for p in all_params: p.grad = None
         return free_ce, beta_t, GOV['K'], False
+    GOV['drift'] = drift
     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
@@ -132,10 +141,11 @@ def ep_step(x, y):
         if g is not None: gn += float((g ** 2).sum())
     gn = gn ** 0.5
     if GOV['gema'] is None: GOV['gema'] = gn
-    if not math.isfinite(gn) or gn > 8 * GOV['gema']:
+    GOV['gema'] = 0.99 * GOV['gema'] + 0.01 * gn        # EMA always updates (frozen-ref bugfix)
+    GOV['gn'] = gn
+    if not math.isfinite(gn) or (gn > 8 * GOV['gema'] and not args.noguard):
         for p in all_params: p.grad = None
         return free_ce, beta_t, GOV['K'], False
-    GOV['gema'] = 0.99 * GOV['gema'] + 0.01 * gn
     for p, g in zip(all_params, gs):
         p.grad = g
     return free_ce, beta_t, GOV['K'], True
@@ -193,7 +203,8 @@ for step in range(args.steps + 1):
         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'| beta={beta_t:.2e} K={rounds} skips={skips}{gtag} | {step/max(time.time()-t0,1e-9):.3f} it/s', flush=True)
+              f'| beta={beta_t:.2e} K={rounds} skips={skips}{gtag} '
+              f'drift={GOV["drift"]:.3f} gn={GOV["gn"]:.2e} sig={GOV["sig"]:.1f} | {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, 'beta_t': beta_t,
                          'rounds': rounds, 'skips': skips, 'gate_cos': (None if math.isnan(gcos) else gcos)}, step=step)