fidelity factorial arms (randk / noises0 / hidden-capacity) + pre-registration item 15

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

scripts/eval_loop_code.py

@@ -93,6 +93,8 @@ def main():
                     help="PerDepthAdapter (Bae-style per-iteration merges)")
     ap.add_argument("--tiedalpha", action="store_true",
                     help="TiedAlphaAdapter (learned per-dim alpha, tied B)")
+    ap.add_argument("--noises0", action="store_true")
+    ap.add_argument("--hidden", type=int, default=512)
     ap.add_argument("--halt", action="store_true",
                     help="record per-item convergence depth (free-ACT probe)")
     args = ap.parse_args()
@@ -101,13 +103,17 @@ def main():
     model, tok = load_model(dtype=torch.bfloat16)
     tok.padding_side = "left"
     looper = BandLooper(model)
-    from loop_common import PerDepthAdapter, TiedAlphaAdapter
-    cls = (TiedAlphaAdapter if args.tiedalpha
+    from loop_common import (NoisyMergeAdapter, PerDepthAdapter,
+                             TiedAlphaAdapter)
+    cls = (NoisyMergeAdapter if args.noises0
+           else TiedAlphaAdapter if args.tiedalpha
            else PerDepthAdapter if args.perdepth
            else 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})
+    if not args.adaptive:
+        kw["hidden"] = args.hidden
     adapter = cls(d=model.config.get_text_config().hidden_size, **kw).cuda()
     if args.adapter:
         adapter.load_state_dict(torch.load(args.adapter, map_location="cuda"))

scripts/loop_common.py

@@ -180,6 +180,17 @@ class ParcaeAdapter(RecurrentAdapter):
             return self.A_diag().max().item()
 
 
+class NoisyMergeAdapter(MergeAdapter):
+    """MergeAdapter with noise s0 (factorial cell: isolates the initial
+    state; everything else identical to the standard merge)."""
+
+    def init_state(self, e):
+        n = torch.randn_like(e.float())
+        n = n * (e.float().norm(dim=-1, keepdim=True)
+                 / (n.norm(dim=-1, keepdim=True) + 1e-6))
+        return n.to(e.dtype)
+
+
 class TiedAlphaAdapter(nn.Module):
     """Merge with LEARNED per-dim constant alpha, B tied to (1-a):
         x = (1-a)*e + a*s_hat + MLP([e;s_hat]),  a = sigmoid(a_hat)

scripts/train_merge_code.py

@@ -21,8 +21,8 @@ import torch
 import torch.nn.functional as F
 
 from loop_common import (AdaptiveMergeAdapter, BandLooper, MergeAdapter,
-                         ParcaeAdapter, PerDepthAdapter, RecurrentAdapter,
-                         TiedAlphaAdapter)
+                         NoisyMergeAdapter, ParcaeAdapter, PerDepthAdapter,
+                         RecurrentAdapter, TiedAlphaAdapter)
 from prep_mbpp import DIRECT_SUFFIX, mbpp_prompt
 
 sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
@@ -62,6 +62,13 @@ 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)")
+ap.add_argument("--randk", action="store_true",
+                help="factorial cell: standard MergeAdapter but rec-style "
+                     "log-uniform random depth (isolates the curriculum)")
+ap.add_argument("--noises0", action="store_true",
+                help="factorial cell: standard merge + noise s0")
+ap.add_argument("--hidden", type=int, default=512,
+                help="adapter MLP hidden width (capacity control)")
 ap.add_argument("--tiedalpha", action="store_true",
                 help="merge with learned per-dim constant alpha, B tied "
                      "to (1-a); standard curriculum")
@@ -71,7 +78,7 @@ ap.add_argument("--perdepth", type=int, default=0,
 ARGS, _ = ap.parse_known_args()
 if ARGS.parcae:
     ARGS.rec = True
-if ARGS.rec and not ARGS.bptt:
+if (ARGS.rec or ARGS.randk) and not ARGS.bptt:
     ARGS.bptt = 4
 SEED = ARGS.seed
 LR = ARGS.lr
@@ -85,7 +92,10 @@ SUFFIX = ((f"_s{SEED}" if SEED else "")
           + (("_parcae" if ARGS.parcae else "_rec") + str(ARGS.recmax)
              if ARGS.rec else "")
           + (f"_pd{ARGS.perdepth}" if ARGS.perdepth else "")
-          + ("_ta" if ARGS.tiedalpha else ""))
+          + ("_ta" if ARGS.tiedalpha else "")
+          + (f"_rk{ARGS.recmax}" if ARGS.randk else "")
+          + ("_ns" if ARGS.noises0 else "")
+          + (f"_h{ARGS.hidden}" if ARGS.hidden != 512 else ""))
 PAUSE_ID = 6  # <unused0>
 
 
@@ -144,8 +154,11 @@ def main():
         p.requires_grad_(False)
     looper = BandLooper(model)
     d = model.config.get_text_config().hidden_size
+    H = ARGS.hidden
     if ARGS.tiedalpha:
-        adapter = TiedAlphaAdapter(d=d, alpha=ARGS.alpha).cuda()
+        adapter = TiedAlphaAdapter(d=d, alpha=ARGS.alpha, hidden=H).cuda()
+    elif ARGS.noises0:
+        adapter = NoisyMergeAdapter(d=d, alpha=ARGS.alpha, hidden=H).cuda()
     elif ARGS.perdepth:
         adapter = PerDepthAdapter(d=d, alpha=ARGS.alpha,
                                   n_depth=ARGS.perdepth).cuda()
@@ -156,7 +169,7 @@ def main():
     elif ARGS.adaptive:
         adapter = AdaptiveMergeAdapter(d=d, alpha0=ARGS.alpha).cuda()
     else:
-        adapter = MergeAdapter(d=d, alpha=ARGS.alpha).cuda()
+        adapter = MergeAdapter(d=d, alpha=ARGS.alpha, hidden=H).cuda()
     if ARGS.warm:
         adapter.load_state_dict(torch.load(ARGS.warm, map_location="cuda"))
         print("warm-started from", ARGS.warm, flush=True)
@@ -183,7 +196,7 @@ def main():
     log = []
     t0 = time.time()
     for step in range(STEPS):
-        if ARGS.rec:
+        if ARGS.rec or ARGS.randk:
             # 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))))))