RecurrentAdapter arm: Huginn-regime retrofit (learned A/B, noise h0, randomized depth) + pre-registration item 11

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

scripts/train_merge_code.py

@@ -20,7 +20,8 @@ from pathlib import Path
 import torch
 import torch.nn.functional as F
 
-from loop_common import AdaptiveMergeAdapter, BandLooper, MergeAdapter
+from loop_common import (AdaptiveMergeAdapter, BandLooper, MergeAdapter,
+                         RecurrentAdapter)
 from prep_mbpp import DIRECT_SUFFIX, mbpp_prompt
 
 sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
@@ -52,7 +53,14 @@ ap.add_argument("--bptt", type=int, default=0,
                 help="truncated BPTT: grads only through last N iterations")
 ap.add_argument("--lr", type=float, default=1e-3)
 ap.add_argument("--warm", default=None, help="warm-start adapter checkpoint")
+ap.add_argument("--rec", action="store_true",
+                help="Huginn-style regime: RecurrentAdapter (learned A/B, "
+                     "noise h0) + log-uniform random depth 1..recmax, "
+                     "truncated bptt (default 4)")
+ap.add_argument("--recmax", type=int, default=16)
 ARGS, _ = ap.parse_known_args()
+if ARGS.rec and not ARGS.bptt:
+    ARGS.bptt = 4
 SEED = ARGS.seed
 LR = ARGS.lr
 SUFFIX = ((f"_s{SEED}" if SEED else "")
@@ -61,7 +69,8 @@ SUFFIX = ((f"_s{SEED}" if SEED else "")
           + ("_ad" if ARGS.adaptive else "")
           + (f"_dk{ARGS.deepk}" if ARGS.deepk else "")
           + (f"_lr{ARGS.lr}" if ARGS.lr != 1e-3 else "")
-          + ("_warm" if ARGS.warm else ""))
+          + ("_warm" if ARGS.warm else "")
+          + (f"_rec{ARGS.recmax}" if ARGS.rec else ""))
 PAUSE_ID = 6  # <unused0>
 
 
@@ -120,7 +129,9 @@ def main():
         p.requires_grad_(False)
     looper = BandLooper(model)
     d = model.config.get_text_config().hidden_size
-    if ARGS.adaptive:
+    if ARGS.rec:
+        adapter = RecurrentAdapter(d=d, alpha=ARGS.alpha).cuda()
+    elif ARGS.adaptive:
         adapter = AdaptiveMergeAdapter(d=d, alpha0=ARGS.alpha).cuda()
     else:
         adapter = MergeAdapter(d=d, alpha=ARGS.alpha).cuda()
@@ -150,7 +161,13 @@ def main():
     log = []
     t0 = time.time()
     for step in range(STEPS):
-        k, labels = K_BUCKETS[step % len(K_BUCKETS)]
+        if ARGS.rec:
+            # randomized depth, log-uniform in [1, recmax], any difficulty
+            k = min(ARGS.recmax,
+                    max(1, int(math.exp(rng.uniform(0, math.log(ARGS.recmax))))))
+            labels = ("easy", "hard")
+        else:
+            k, labels = K_BUCKETS[step % len(K_BUCKETS)]
         cand = [it for lbl in labels for it in pool[lbl]]
         batch = rng.sample(cand, min(BATCH, len(cand)))
         ids, msk, lab, lmask = build_code_batch(tok, batch)
@@ -174,7 +191,7 @@ def main():
                   f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True)
         if step % 100 == 99 or step == STEPS - 1:
             vals = {}
-            for kk in (0, 1, 2, 4):
+            for kk in ((0, 1, 2, 4, 8, 16) if ARGS.rec else (0, 1, 2, 4)):
                 vals[f"easy_k{kk}"] = val_loss(looper, adapter, tok,
                                                val["easy"], kk)
                 vals[f"hard_k{kk}"] = val_loss(looper, adapter, tok,