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/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)):