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/eval_loop_code.py

@@ -15,7 +15,8 @@ from pathlib import Path
 
 import torch
 
-from loop_common import AdaptiveMergeAdapter, BandLooper, MergeAdapter
+from loop_common import (AdaptiveMergeAdapter, BandLooper, MergeAdapter,
+                         RecurrentAdapter)
 from prep_mbpp import DIRECT_SUFFIX, extract_code, mbpp_prompt, run_tests
 
 sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
@@ -76,13 +77,16 @@ def main():
                     help="append p pause tokens to each prompt")
     ap.add_argument("--alpha", type=float, default=0.3)
     ap.add_argument("--adaptive", action="store_true")
+    ap.add_argument("--rec", action="store_true",
+                    help="RecurrentAdapter (noise h0, learned A/B)")
     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 = AdaptiveMergeAdapter if args.adaptive else MergeAdapter
+    cls = (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()
     if args.adapter:

scripts/loop_common.py

@@ -84,6 +84,50 @@ class AdaptiveMergeAdapter(nn.Module):
         return out.to(dt)
 
 
+class RecurrentAdapter(nn.Module):
+    """Huginn-style recurrent-state update on the frozen band
+    (arXiv 2502.05171 regime: h_{t+1} = A·h_t + B·e + Transformer(h_t, e)).
+
+    The band's residual stream supplies the "+Transformer" term, so the
+    adapter computes the band input x_t = A·ĥ_t + B·e + MLP([e;ĥ_t]) with
+    LEARNED d×d maps A, B (init A=α·I, B=(1−α)·I: starts exactly at the
+    fixed merge). h_0 is norm-scaled noise via init_state — combined with
+    randomized-depth training this targets depth-monotone iteration rather
+    than our anchor-dominant fixed point.
+    """
+
+    def __init__(self, d=1536, hidden=512, alpha=ALPHA, sigma=1.0):
+        super().__init__()
+        self.A = nn.Linear(d, d, bias=False)
+        self.B = nn.Linear(d, d, bias=False)
+        with torch.no_grad():
+            self.A.weight.copy_(alpha * torch.eye(d))
+            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 init_state(self, e):
+        """h_0: per-position Gaussian noise scaled to the anchor's norm."""
+        n = torch.randn_like(e.float())
+        n = n * (e.float().norm(dim=-1, keepdim=True)
+                 / (n.norm(dim=-1, keepdim=True) + 1e-6)) * self.sigma
+        return n.to(e.dtype)
+
+    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(s_hat) + self.B(e32)
+               + self.mlp(torch.cat([e32, s_hat], dim=-1)))
+        return out.to(dt)
+
+
 class BandLooper:
     """Capture layer-call kwargs once per forward, then re-run L14-30 manually."""
 
@@ -181,7 +225,10 @@ class BandLooper:
             logits = self.suffix_logits(s, calls, last_only=last_only)
             return (logits, [s]) if return_states else logits
         with torch.no_grad():
-            s = self.band(e, calls)  # s_0: no trainable params upstream
+            # s_0: no trainable params upstream (noise state for recurrent
+            # adapters — the band(e) warm start would hide the B·e path)
+            s = (adapter.init_state(e) if hasattr(adapter, "init_state")
+                 else self.band(e, calls))
         states = [s]
         n_nograd = max(0, k - bptt) if bptt else 0
         for i in range(k):
@@ -270,7 +317,8 @@ class BandLooper:
             calls, _ = self.capture(input_ids, attention_mask,
                                     logits_to_keep=1)
             e = self._hin[self.l0]
-            s = self.band(e, calls)
+            s = (adapter.init_state(e) if hasattr(adapter, "init_state")
+                 else self.band(e, calls))
             x_star = e
             for _ in range(k):
                 x_star = adapter(e, s)

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,