From d00f120f27142985a5dc93e45d4c4d22cabf5432 Mon Sep 17 00:00:00 2001 From: Nils Date: Wed, 15 Jul 2026 01:07:01 +0200 Subject: [PATCH] PerDepthAdapter (Bae-style per-iteration merges) + convergence-halting probe (free ACT); pre-registration item 13 Co-Authored-By: Claude Fable 5 --- results-loop/PROTOCOL_UNIFIED.md | 16 ++++++++++++ scripts/eval_loop_code.py | 28 +++++++++++++++++--- scripts/loop_common.py | 44 +++++++++++++++++++++++++++----- scripts/train_merge_code.py | 13 +++++++--- 4 files changed, 88 insertions(+), 13 deletions(-) diff --git a/results-loop/PROTOCOL_UNIFIED.md b/results-loop/PROTOCOL_UNIFIED.md index ecbbd04..a50d605 100644 --- a/results-loop/PROTOCOL_UNIFIED.md +++ b/results-loop/PROTOCOL_UNIFIED.md @@ -133,3 +133,19 @@ number for the unified adapter exists at time of writing. stability is where genuine iteration must live. Either outcome formalizes "the anchor coefficient is the stability dial" as "the anchor coefficient is the spectral radius". + +13. **Per-depth adapter arm + free-ACT probe (pre-registered 2026-07-15, + before training).** (a) PerDepthAdapter: one merge adapter per + iteration (n=4, Bae-style depth-wise relaxation at the entrance; + breaks time-invariance — LTV, no fixed-point guarantee), standard + curriculum, e400, eval ks 0,2,4,8. Prediction: lands at or below the + distill/rung-2 amortization ceiling (~46% hard) because depth-indexed + weights add content, not state-evolution; exceeding it would show + per-iteration expressivity was binding and amend the deflationary + claim. Depths >4 reuse adapter 4 (stated: k=8 cell is then + fixed-point-like by construction). (b) Free-ACT probe on the standard + merge arm: record per-item convergence depth (cos>0.9995) at k=8 cap. + Predictions: accuracy unchanged vs fixed k (post-convergence no-ops); + mean k_conv ≈ 3; hard-labeled items converge SLOWER than easy ones + (adaptive compute allocates like ACT without any learned halting + parameter). diff --git a/scripts/eval_loop_code.py b/scripts/eval_loop_code.py index 9f1f958..78b3197 100644 --- a/scripts/eval_loop_code.py +++ b/scripts/eval_loop_code.py @@ -28,8 +28,9 @@ OUT = Path(os.environ.get("LOOP_OUT", @torch.no_grad() def pass1_at_k(looper, adapter, tok, items, k, batch=8, max_new=220, - feedforward=False, pause=0): + feedforward=False, pause=0, halt=False): codes = [] + k_convs = [] for i in range(0, len(items), batch): torch.cuda.empty_cache() chunk = items[i : i + batch] @@ -42,10 +43,14 @@ def pass1_at_k(looper, adapter, tok, items, k, batch=8, max_new=220, enc["input_ids"] = torch.cat([enc["input_ids"], pcol], 1) enc["attention_mask"] = torch.cat( [enc["attention_mask"], torch.ones_like(pcol)], 1) + cv = {} if halt else None gen = looper.generate_frozen_prompt(adapter, tok, enc["input_ids"], k, max_new_tokens=max_new, attention_mask=enc["attention_mask"], - feedforward=feedforward) + feedforward=feedforward, + conv_out=cv) + if halt: + k_convs.extend(cv.get("k_conv", [])) for j in range(len(chunk)): txt = tok.decode(gen[j, enc["input_ids"].shape[1]:], skip_special_tokens=True) @@ -61,6 +66,9 @@ def pass1_at_k(looper, adapter, tok, items, k, batch=8, max_new=220, d[1] += 1 per_item = [{"task_id": it["task_id"], "ok": bool(ok)} for it, ok in zip(items, oks)] + if halt and k_convs: + for r, kc in zip(per_item, k_convs): + r["k_conv"] = kc return (hits / len(items), {l: c / n for l, (c, n) in per_label.items()}, per_item) @@ -81,13 +89,19 @@ def main(): help="RecurrentAdapter (noise h0, learned A/B)") ap.add_argument("--parcae", action="store_true", help="ParcaeAdapter (rho(A)<1 by construction)") + ap.add_argument("--perdepth", action="store_true", + help="PerDepthAdapter (Bae-style per-iteration merges)") + ap.add_argument("--halt", action="store_true", + help="record per-item convergence depth (free-ACT probe)") 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 = (ParcaeAdapter if args.parcae + from loop_common import PerDepthAdapter + cls = (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}) @@ -106,7 +120,13 @@ def main(): t0 = time.time() acc, by_label, per_item = pass1_at_k(looper, adapter, tok, items, k, feedforward=args.feedforward, - pause=args.pause) + pause=args.pause, halt=args.halt) + if args.halt: + by_lbl_k = {} + for it, r in zip(items, per_item): + by_lbl_k.setdefault(it["label"], []).append(r.get("k_conv", k)) + print(" mean k_conv:", {l: round(sum(v)/len(v), 2) + for l, v in by_lbl_k.items()}, flush=True) res["ks"][k] = {"acc": acc, "by_label": by_label, "per_item": per_item} print(f"k={k}: pass@1={acc:.3f} " diff --git a/scripts/loop_common.py b/scripts/loop_common.py index ad541e0..523c171 100644 --- a/scripts/loop_common.py +++ b/scripts/loop_common.py @@ -180,6 +180,25 @@ class ParcaeAdapter(RecurrentAdapter): return self.A_diag().max().item() +class PerDepthAdapter(nn.Module): + """Depth-wise relaxation (Bae et al. 2024, entrance-level): iteration t + gets its OWN merge adapter — breaks time-invariance, so each loop step + can perform a different computation phase instead of converging to a + fixed point. Depths beyond n_depth reuse the last adapter.""" + + def __init__(self, d=1536, hidden=512, alpha=ALPHA, n_depth=4): + super().__init__() + self.steps = nn.ModuleList( + MergeAdapter(d=d, hidden=hidden, alpha=alpha) + for _ in range(n_depth)) + + def at(self, t): + return self.steps[min(t, len(self.steps) - 1)] + + def forward(self, e, s): # fallback: first-depth adapter + return self.steps[0](e, s) + + class BandLooper: """Capture layer-call kwargs once per forward, then re-run L14-30 manually.""" @@ -286,14 +305,15 @@ class BandLooper: for i in range(k): if i < n_nograd: with torch.no_grad(): - x = adapter(e, s) + x = (adapter.at(i) if hasattr(adapter, "at") + else adapter)(e, s) if loop_mask is not None: x = torch.where(loop_mask[..., None], x, e) s = self.band(x, calls) s = s.detach() states.append(s) continue - x = adapter(e, s) + x = (adapter.at(i) if hasattr(adapter, "at") else adapter)(e, s) if loop_mask is not None: x = torch.where(loop_mask[..., None], x, e) if use_checkpoint: @@ -351,7 +371,8 @@ class BandLooper: @torch.no_grad() def generate_frozen_prompt(self, adapter, tok, input_ids, k, max_new_tokens=220, attention_mask=None, - stop_strs=(), feedforward=False): + stop_strs=(), feedforward=False, + conv_out=None): """Fast equivalent of loop_generate(loop_prompt_only=True). The looped prompt states are constant across token steps (causality), @@ -372,9 +393,20 @@ class BandLooper: 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) - s = self.band(x_star, calls) + conv = (torch.full((e.shape[0],), -1, dtype=torch.long) + if conv_out is not None else None) + for _i in range(k): + x_star = (adapter.at(_i) if hasattr(adapter, "at") + else adapter)(e, s) + s_new = self.band(x_star, calls) + if conv is not None: + c = nn.functional.cosine_similarity( + s_new.float().flatten(1), s.float().flatten(1), dim=1) + conv[((c > 0.9995).cpu()) & (conv < 0)] = _i + 1 + s = s_new + if conv is not None: + conv[conv < 0] = k + conv_out["k_conv"] = conv.tolist() del calls # prefill re-runs band(x_star) -> same final s as slow path hook = None diff --git a/scripts/train_merge_code.py b/scripts/train_merge_code.py index 68602c8..b642e0a 100644 --- a/scripts/train_merge_code.py +++ b/scripts/train_merge_code.py @@ -21,7 +21,7 @@ import torch import torch.nn.functional as F from loop_common import (AdaptiveMergeAdapter, BandLooper, MergeAdapter, - ParcaeAdapter, RecurrentAdapter) + ParcaeAdapter, PerDepthAdapter, RecurrentAdapter) from prep_mbpp import DIRECT_SUFFIX, mbpp_prompt sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) @@ -61,6 +61,9 @@ 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("--perdepth", type=int, default=0, + help="PerDepthAdapter: one merge per iteration depth " + "(Bae-style relaxation), standard curriculum") ARGS, _ = ap.parse_known_args() if ARGS.parcae: ARGS.rec = True @@ -76,7 +79,8 @@ SUFFIX = ((f"_s{SEED}" if SEED else "") + (f"_lr{ARGS.lr}" if ARGS.lr != 1e-3 else "") + ("_warm" if ARGS.warm else "") + (("_parcae" if ARGS.parcae else "_rec") + str(ARGS.recmax) - if ARGS.rec else "")) + if ARGS.rec else "") + + (f"_pd{ARGS.perdepth}" if ARGS.perdepth else "")) PAUSE_ID = 6 # @@ -135,7 +139,10 @@ def main(): p.requires_grad_(False) looper = BandLooper(model) d = model.config.get_text_config().hidden_size - if ARGS.parcae: + if ARGS.perdepth: + adapter = PerDepthAdapter(d=d, alpha=ARGS.alpha, + n_depth=ARGS.perdepth).cuda() + elif ARGS.parcae: adapter = ParcaeAdapter(d=d, alpha=ARGS.alpha).cuda() elif ARGS.rec: adapter = RecurrentAdapter(d=d, alpha=ARGS.alpha).cuda()