ParcaeAdapter: rho(A)<1 by construction (ZOH negative-diag), rho logging in rec arms, pre-registration item 12

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

scripts/eval_loop_code.py

@@ -16,7 +16,7 @@ from pathlib import Path
 import torch
 
 from loop_common import (AdaptiveMergeAdapter, BandLooper, MergeAdapter,
-                         RecurrentAdapter)
+                         ParcaeAdapter, RecurrentAdapter)
 from prep_mbpp import DIRECT_SUFFIX, extract_code, mbpp_prompt, run_tests
 
 sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
@@ -79,13 +79,16 @@ def main():
     ap.add_argument("--adaptive", action="store_true")
     ap.add_argument("--rec", action="store_true",
                     help="RecurrentAdapter (noise h0, learned A/B)")
+    ap.add_argument("--parcae", action="store_true",
+                    help="ParcaeAdapter (rho(A)<1 by construction)")
     args = ap.parse_args()
     ks = [int(x) for x in args.ks.split(",")]
 
     model, tok = load_model(dtype=torch.bfloat16)
     tok.padding_side = "left"
     looper = BandLooper(model)
-    cls = (RecurrentAdapter if args.rec
+    cls = (ParcaeAdapter if args.parcae
+           else RecurrentAdapter if args.rec
            else AdaptiveMergeAdapter if args.adaptive else MergeAdapter)
     kw = ({"alpha0": args.alpha} if args.adaptive else {"alpha": args.alpha})
     adapter = cls(d=model.config.get_text_config().hidden_size, **kw).cuda()

scripts/loop_common.py

@@ -127,6 +127,58 @@ class RecurrentAdapter(nn.Module):
                + self.mlp(torch.cat([e32, s_hat], dim=-1)))
         return out.to(dt)
 
+    def rho(self):
+        """Spectral radius of A (power iteration) — the stability dial."""
+        with torch.no_grad():
+            v = torch.randn(self.A.weight.shape[0],
+                            device=self.A.weight.device)
+            for _ in range(30):
+                v = self.A.weight.T @ (self.A.weight @ v)
+                v = v / (v.norm() + 1e-12)
+            return (self.A.weight @ v).norm().item()
+
+
+class ParcaeAdapter(RecurrentAdapter):
+    """Recurrent adapter with ρ(A) < 1 GUARANTEED by construction
+    (Parcae parameterization, Prairie et al. 2026): A is the ZOH
+    discretization of a continuous negative-diagonal system,
+    A = exp(−Δt·exp(a)) with learned per-dim a and scalar log-Δt, so every
+    entry lies in (0,1) regardless of optimizer noise. Init matches the
+    fixed merge exactly: a = log(−log α), Δt = 1 → A = α·I."""
+
+    def __init__(self, d=1536, hidden=512, alpha=ALPHA, sigma=1.0):
+        nn.Module.__init__(self)
+        import math as _m
+        self.log_a = nn.Parameter(
+            torch.full((d,), _m.log(-_m.log(alpha))))
+        self.log_dt = nn.Parameter(torch.zeros(()))
+        self.B = nn.Linear(d, d, bias=False)
+        with torch.no_grad():
+            self.B.weight.copy_((1 - alpha) * torch.eye(d))
+        self.mlp = nn.Sequential(
+            nn.Linear(2 * d, hidden), nn.GELU(), nn.Linear(hidden, d)
+        )
+        nn.init.zeros_(self.mlp[2].weight)
+        nn.init.zeros_(self.mlp[2].bias)
+        self.sigma = sigma
+
+    def A_diag(self):
+        return torch.exp(-torch.exp(self.log_dt) * torch.exp(self.log_a))
+
+    def forward(self, e, s):
+        dt = e.dtype
+        e32, s32 = e.float(), s.float()
+        s_hat = s32 * (
+            e32.norm(dim=-1, keepdim=True) / (s32.norm(dim=-1, keepdim=True) + 1e-6)
+        )
+        out = (self.A_diag().to(e32.dtype) * s_hat + self.B(e32)
+               + self.mlp(torch.cat([e32, s_hat], dim=-1)))
+        return out.to(dt)
+
+    def rho(self):
+        with torch.no_grad():
+            return self.A_diag().max().item()
+
 
 class BandLooper:
     """Capture layer-call kwargs once per forward, then re-run L14-30 manually."""

scripts/train_merge_code.py

@@ -21,7 +21,7 @@ import torch
 import torch.nn.functional as F
 
 from loop_common import (AdaptiveMergeAdapter, BandLooper, MergeAdapter,
-                         RecurrentAdapter)
+                         ParcaeAdapter, RecurrentAdapter)
 from prep_mbpp import DIRECT_SUFFIX, mbpp_prompt
 
 sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
@@ -58,7 +58,12 @@ ap.add_argument("--rec", action="store_true",
                      "noise h0) + log-uniform random depth 1..recmax, "
                      "truncated bptt (default 4)")
 ap.add_argument("--recmax", type=int, default=16)
+ap.add_argument("--parcae", action="store_true",
+                help="rec regime with rho(A)<1 by construction "
+                     "(diag-negative-exp ZOH parameterization)")
 ARGS, _ = ap.parse_known_args()
+if ARGS.parcae:
+    ARGS.rec = True
 if ARGS.rec and not ARGS.bptt:
     ARGS.bptt = 4
 SEED = ARGS.seed
@@ -70,7 +75,8 @@ SUFFIX = ((f"_s{SEED}" if SEED else "")
           + (f"_dk{ARGS.deepk}" if ARGS.deepk else "")
           + (f"_lr{ARGS.lr}" if ARGS.lr != 1e-3 else "")
           + ("_warm" if ARGS.warm else "")
-          + (f"_rec{ARGS.recmax}" if ARGS.rec else ""))
+          + (("_parcae" if ARGS.parcae else "_rec") + str(ARGS.recmax)
+             if ARGS.rec else ""))
 PAUSE_ID = 6  # <unused0>
 
 
@@ -129,7 +135,9 @@ def main():
         p.requires_grad_(False)
     looper = BandLooper(model)
     d = model.config.get_text_config().hidden_size
-    if ARGS.rec:
+    if ARGS.parcae:
+        adapter = ParcaeAdapter(d=d, alpha=ARGS.alpha).cuda()
+    elif ARGS.rec:
         adapter = RecurrentAdapter(d=d, alpha=ARGS.alpha).cuda()
     elif ARGS.adaptive:
         adapter = AdaptiveMergeAdapter(d=d, alpha0=ARGS.alpha).cuda()
@@ -187,8 +195,12 @@ def main():
 
         log.append({"step": step, "k": k, "loss": loss.item()})
         if step % 10 == 0:
-            print(f"step {step:4d}  k={k}  loss={loss.item():.4f}  "
+            rho = (f"  rho(A)={adapter.rho():.3f}"
+                   if hasattr(adapter, "rho") else "")
+            print(f"step {step:4d}  k={k}  loss={loss.item():.4f}{rho}  "
                   f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True)
+            if rho:
+                log.append({"step": step, "rho": adapter.rho()})
         if step % 100 == 99 or step == STEPS - 1:
             vals = {}
             for kk in ((0, 1, 2, 4, 8, 16) if ARGS.rec else (0, 1, 2, 4)):