diff --git a/results-loop/PROTOCOL_UNIFIED.md b/results-loop/PROTOCOL_UNIFIED.md index d9bf387..52542ac 100644 --- a/results-loop/PROTOCOL_UNIFIED.md +++ b/results-loop/PROTOCOL_UNIFIED.md @@ -101,3 +101,19 @@ number for the unified adapter exists at time of writing. (9,30) overall 14.0-21.4%, hard ≤21.4% — catastrophic, like anchors 11-13, despite L9 being a full-attention KV-computing layer. The lens boundary, not layer type, gates the retrofit. + +11. **Recurrent-regime arm (pre-registered 2026-07-15, before training).** + Huginn-style retrofit on the frozen E2B band: RecurrentAdapter + (learned A,B init α·I/(1−α)·I + zero-init MLP), h0 = norm-scaled + noise, log-uniform random depth k∈[1,16], bptt=4, same data/steps/ + checkpoint rule (e400 primary) as all merge arms. Eval ks 0,2,4,8,16,32 + on the 250-item MBPP set. Competing predictions: (a) "amortization is + intrinsic to frozen-band retrofits" → performance plateaus by k≈4 at + or below the merge arm's level, no depth-monotone gain; (b) "fixed- + point behavior was an artifact of our fixed-shallow-k training" + (Huginn regime transfers) → monotone hard-bucket improvement past k=8 + and reduced noise-seed sensitivity after training. Secondary readout: + path independence (two noise seeds → output agreement rate) at e400. + Known risk, stated in advance: 600 steps may be far too little for + this regime (McLeish et al. use ~50B tokens); a null here bounds the + cheap-retrofit budget only, not the regime. diff --git a/scripts/eval_loop_code.py b/scripts/eval_loop_code.py index 2cb6108..2004955 100644 --- a/scripts/eval_loop_code.py +++ b/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: diff --git a/scripts/loop_common.py b/scripts/loop_common.py index fb46dc6..8976e88 100644 --- a/scripts/loop_common.py +++ b/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) diff --git a/scripts/train_merge_code.py b/scripts/train_merge_code.py index 6631310..a61fc7e 100644 --- a/scripts/train_merge_code.py +++ b/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 # @@ -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,