E1 halting-gate machinery (adapter, trainer, gated eval) + pre-registration item 18

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

results-loop/PROTOCOL_UNIFIED.md

@@ -359,3 +359,21 @@ on GSM alone in the correct regime, delivers no usable gain and the
 standard fidelity damage; combined with the scope note, the boundary
 claim (structural: supervision density + state-evolution bottleneck)
 is fully supported.
+
+18. **E1: learned per-prompt halting gate (pre-registered 2026-07-16
+    ~00:20, before any arm runs; PLAN_SELFPACED.md).** HaltingMergeAdapter:
+    frozen-recipe merge + ACT-style halting head on the last prompt
+    position's workspace state; soft state-mixture training, CE + lambda *
+    E[iters], penalty warmup at step 100; NO difficulty curriculum (mixed
+    batches — the gate must discover the allocation). k_max=4, e400/e600
+    checkpoints, deploy = sequential halting at 0.5 cumulative mass,
+    generation via frozen-prompt at per-item k*. Arms: lambda in
+    {0, 1e-3, 1e-2}, seed 0. Eval: 250 items, vs anchors k=0 (0.488),
+    uniform merge k=4 (0.512/0.885/0.464), probe-gate E0 (0.520/0.975/0.286).
+    Predictions: (a) some lambda gives overall >= 0.512 at mean E[k] <=
+    2.4 (60% of uniform-4); (b) easy >= 0.95 at that lambda; (c) k*-vs-hard
+    point-biserial r > 0.3; (d) hard >= 0.286 (beats E0's frozen probe).
+    Collapse (E[k] pinned at 1 or 4 for all lambda) falsifies E1 and
+    triggers the plan's kill criterion. lambda=0 control isolates whether
+    the CE gradient alone moves the gate (expected: barely — penalty
+    provides the pressure).

scripts/eval_gate_code.py

@@ -0,0 +1,114 @@
+"""E1 eval (PLAN_SELFPACED): halting-gate MBPP eval.
+
+Per item: deploy-time halting picks k* (1..kmax); generation then uses the
+frozen-prompt path at that k (items grouped by k* for batching). Reports
+pass@1 by label, the k* distribution by label, mean compute, and the
+gate-difficulty point-biserial correlation.
+"""
+
+import argparse
+import json
+import os
+import sys
+import time
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+import torch
+
+from halting_common import HaltingMergeAdapter, halted_k_per_item
+from loop_common import BandLooper
+from prep_mbpp import DIRECT_SUFFIX, extract_code, mbpp_prompt, run_tests
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from jlens.core import load_model  # noqa: E402
+
+OUT = Path(os.environ.get("LOOP_OUT",
+    Path(__file__).resolve().parent.parent / "results-loop"))
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--adapter", required=True)
+    ap.add_argument("--tag", required=True)
+    ap.add_argument("--kmax", type=int, default=4)
+    ap.add_argument("--n", type=int, default=250)
+    ap.add_argument("--batch", type=int, default=8)
+    args = ap.parse_args()
+
+    model, tok = load_model(dtype=torch.bfloat16)
+    tok.padding_side = "left"
+    looper = BandLooper(model)
+    adapter = HaltingMergeAdapter(
+        d=model.config.get_text_config().hidden_size).cuda()
+    adapter.load_state_dict(torch.load(args.adapter, map_location="cuda"))
+    adapter.eval()
+
+    items = [it for it in json.load(open(OUT / "mbpp_data.json"))
+             if it["split"] == "test"][: args.n]
+    print(f"[{args.tag}] gated MBPP eval on {len(items)}, kmax={args.kmax}",
+          flush=True)
+
+    # phase 1: per-item k*
+    kstars = []
+    with torch.no_grad():
+        for i in range(0, len(items), args.batch):
+            chunk = items[i : i + args.batch]
+            enc = tok([mbpp_prompt(tok, it, DIRECT_SUFFIX) for it in chunk],
+                      return_tensors="pt", padding=True,
+                      add_special_tokens=False).to("cuda")
+            pl = enc["attention_mask"].sum(-1) - 1  # left-pad: last position
+            pl = torch.full_like(pl, enc["input_ids"].shape[1] - 1)
+            ks = halted_k_per_item(looper, adapter, enc["input_ids"],
+                                   args.kmax, enc["attention_mask"], pl)
+            kstars.extend(ks.tolist())
+    t0 = time.time()
+
+    # phase 2: generate grouped by k*
+    codes = [None] * len(items)
+    for kval in sorted(set(kstars)):
+        idxs = [i for i, kk in enumerate(kstars) if kk == kval]
+        for j in range(0, len(idxs), args.batch):
+            grp = idxs[j : j + args.batch]
+            enc = tok([mbpp_prompt(tok, items[i], DIRECT_SUFFIX) for i in grp],
+                      return_tensors="pt", padding=True,
+                      add_special_tokens=False).to("cuda")
+            gen = looper.generate_frozen_prompt(
+                adapter, tok, enc["input_ids"], kval, max_new_tokens=220,
+                attention_mask=enc["attention_mask"])
+            for gi, i in enumerate(grp):
+                txt = tok.decode(gen[gi, enc["input_ids"].shape[1]:],
+                                 skip_special_tokens=True)
+                codes[i] = extract_code(txt)
+    with ThreadPoolExecutor(8) as ex:
+        oks = list(ex.map(lambda ci: run_tests(ci[0], ci[1]),
+                          zip(codes, items)))
+
+    per_label, kdist = {}, {}
+    for it, ok, kk in zip(items, oks, kstars):
+        d = per_label.setdefault(it["label"], [0, 0, 0.0])
+        d[0] += ok; d[1] += 1; d[2] += kk
+        kdist.setdefault(it["label"], []).append(kk)
+    acc = sum(oks) / len(items)
+    by_label = {l: c / n for l, (c, n, _) in per_label.items()}
+    mean_k = {l: sum(v) / len(v) for l, v in kdist.items()}
+    hard = torch.tensor([it["label"] == "hard" for it in items], dtype=torch.float)
+    kk = torch.tensor(kstars, dtype=torch.float)
+    r = ((kk - kk.mean()) * (hard - hard.mean())).mean() / (kk.std() * hard.std() + 1e-9)
+    print(f"pass@1={acc:.3f} by_label={ {l: round(v,3) for l,v in by_label.items()} }")
+    print(f"mean k* by label: { {l: round(v,2) for l,v in mean_k.items()} }  "
+          f"overall E[k]={sum(kstars)/len(kstars):.2f}  "
+          f"gate-difficulty r={r:.3f}  ({time.time()-t0:.0f}s)", flush=True)
+
+    json.dump({"tag": args.tag, "acc": acc, "by_label": by_label,
+               "mean_kstar": mean_k, "corr_hard": r.item(),
+               "kstars": kstars,
+               "per_item": [{"task_id": it["task_id"], "ok": bool(o),
+                             "kstar": kk_} for it, o, kk_ in
+                            zip(items, oks, kstars)]},
+              open(OUT / f"eval_code_{args.tag}.json", "w"), indent=1)
+    print("wrote", OUT / f"eval_code_{args.tag}.json")
+
+
+if __name__ == "__main__":
+    main()

scripts/halting_common.py

@@ -0,0 +1,101 @@
+"""E1 (PLAN_SELFPACED): learned per-prompt halting over the workspace loop.
+
+HaltingMergeAdapter = frozen-recipe MergeAdapter + a halting head read off
+the workspace state at the last prompt position after each iteration:
+
+    p_i = sigmoid(w · [e_last ; s_hat_i,last] + b)        (halt after iter i)
+    q_i = p_i * prod_{j<i}(1 - p_j),  q_kmax += remainder (ACT-style)
+
+Training (soft, no RL): s_mix = sum_i q_i * s_i  -> suffix -> CE, plus a
+compute penalty  lambda * E[iters] = lambda * sum_i q_i * i.
+The halting head's weight is zero-init and its bias starts at -6
+(p ~ 0.0025), so at init >=98% of the mass sits on k_max: the model is
+approximately the fixed-k merge (documented tolerance, not bit-exact).
+
+Deploy: iterate until cumulative halt mass crosses 0.5; per-item k*.
+"""
+
+import torch
+import torch.nn as nn
+from torch.utils.checkpoint import checkpoint
+
+from loop_common import ALPHA, MergeAdapter
+
+
+class HaltingMergeAdapter(MergeAdapter):
+
+    def __init__(self, d=1536, hidden=512, alpha=ALPHA, bias0=-6.0):
+        super().__init__(d=d, hidden=hidden, alpha=alpha)
+        self.halt = nn.Linear(2 * d, 1)
+        nn.init.zeros_(self.halt.weight)
+        nn.init.constant_(self.halt.bias, bias0)
+
+    def halt_prob(self, e_last, s_last):
+        """p(halt | state) from the last prompt position. (B,) in (0,1)."""
+        e32, s32 = e_last.float(), s_last.float()
+        s_hat = s32 * (e32.norm(dim=-1, keepdim=True)
+                       / (s32.norm(dim=-1, keepdim=True) + 1e-6))
+        return torch.sigmoid(self.halt(torch.cat([e32, s_hat], -1))).squeeze(-1)
+
+
+def halting_loop_logits(looper, adapter, input_ids, k_max, attention_mask,
+                        prompt_last_idx, loop_mask=None, use_checkpoint=True):
+    """Soft-halting teacher-forced logits.
+
+    prompt_last_idx: (B,) index of each item's last prompt position.
+    Returns (logits_of_mixed_state, q, exp_iters):
+      q (B, k_max) halting distribution over iterations 1..k_max.
+    """
+    calls, _ = looper.capture(input_ids, attention_mask, logits_to_keep=1)
+    e = looper._hin[looper.l0].detach()
+    B = e.shape[0]
+    bidx = torch.arange(B, device=e.device)
+    with torch.no_grad():
+        s = looper.band(e, calls)
+    s_mix = torch.zeros_like(e)
+    keep = torch.ones(B, device=e.device)          # prob of not-yet-halted
+    qs = []
+    for i in range(k_max):
+        x = adapter(e, s)
+        if loop_mask is not None:
+            x = torch.where(loop_mask[..., None], x, e)
+        if use_checkpoint:
+            s = checkpoint(lambda x_: looper.band(x_, calls), x,
+                           use_reentrant=False)
+        else:
+            s = looper.band(x, calls)
+        p = adapter.halt_prob(e[bidx, prompt_last_idx],
+                              s[bidx, prompt_last_idx])
+        q_i = keep * p if i < k_max - 1 else keep   # remainder to k_max
+        keep = keep * (1 - p)
+        s_mix = s_mix + q_i[:, None, None] * s.float()
+        qs.append(q_i)
+    q = torch.stack(qs, -1)                          # (B, k_max)
+    exp_iters = (q * torch.arange(1, k_max + 1, device=q.device)).sum(-1)
+    logits = looper.suffix_logits(s_mix.to(e.dtype), calls)
+    return logits, q, exp_iters
+
+
+@torch.no_grad()
+def halted_k_per_item(looper, adapter, input_ids, k_max, attention_mask,
+                      prompt_last_idx, thresh=0.5):
+    """Deploy-time halting: smallest k where cumulative halt mass >= thresh
+    (k_max if never). Returns (B,) ints in 1..k_max."""
+    calls, _ = looper.capture(input_ids, attention_mask, logits_to_keep=1)
+    e = looper._hin[looper.l0].detach()
+    B = e.shape[0]
+    bidx = torch.arange(B, device=e.device)
+    s = looper.band(e, calls)
+    keep = torch.ones(B, device=e.device)
+    cum = torch.zeros(B, device=e.device)
+    kstar = torch.full((B,), k_max, dtype=torch.long, device=e.device)
+    for i in range(k_max):
+        x = adapter(e, s)
+        s = looper.band(x, calls)
+        p = adapter.halt_prob(e[bidx, prompt_last_idx],
+                              s[bidx, prompt_last_idx])
+        cum = cum + keep * p
+        keep = keep * (1 - p)
+        newly = (cum >= thresh) & (kstar == k_max) & (i < k_max - 1)
+        kstar[newly] = i + 1
+    return kstar

scripts/train_gate_code.py

@@ -0,0 +1,120 @@
+"""E1 trainer (PLAN_SELFPACED): halting-gate merge on MBPP prompt-only loop.
+
+Same data/recipe as train_merge_code.py (STaR-verified code supervision,
+prompt-only loop_mask), but NO difficulty->depth curriculum: every batch
+mixes difficulties and runs the soft-halting loop to k_max — the GATE must
+learn the allocation the curriculum used to hand-code.
+
+Loss: CE(code tokens | s_mix) + lam * mean(E[iters]).
+"""
+
+import argparse
+import json
+import os
+import math
+import random
+import sys
+import time
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+from halting_common import HaltingMergeAdapter, halting_loop_logits
+from loop_common import BandLooper
+from prep_mbpp import DIRECT_SUFFIX, mbpp_prompt
+from train_merge_code import build_code_batch, MAX_TOK
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from jlens.core import load_model  # noqa: E402
+
+OUT = Path(os.environ.get("LOOP_OUT",
+    Path(__file__).resolve().parent.parent / "results-loop"))
+STEPS = 600
+BATCH = 4
+WARMUP = 20
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--seed", type=int, default=0)
+ap.add_argument("--lam", type=float, default=1e-3,
+                help="compute penalty weight on E[iterations]")
+ap.add_argument("--kmax", type=int, default=4)
+ap.add_argument("--lr", type=float, default=1e-3)
+ap.add_argument("--lam-warmup", type=int, default=100,
+                help="steps before the penalty ramps in (anti-collapse)")
+ARGS = ap.parse_args()
+LR = ARGS.lr
+TAG = (f"gate_l{ARGS.lam:g}" + (f"_s{ARGS.seed}" if ARGS.seed else ""))
+
+
+def lr_at(step):
+    if step < WARMUP:
+        return LR * (step + 1) / WARMUP
+    t = (step - WARMUP) / max(1, STEPS - WARMUP)
+    return 1e-4 + 0.5 * (LR - 1e-4) * (1 + math.cos(math.pi * t))
+
+
+def lam_at(step):
+    return ARGS.lam * min(1.0, max(0.0, (step - ARGS.lam_warmup) / 100))
+
+
+def main():
+    rng = random.Random(ARGS.seed)
+    torch.manual_seed(ARGS.seed)
+    data = json.load(open(OUT / "mbpp_data.json"))
+    model, tok = load_model(dtype=torch.bfloat16)
+    for p in model.parameters():
+        p.requires_grad_(False)
+    looper = BandLooper(model)
+    adapter = HaltingMergeAdapter(
+        d=model.config.get_text_config().hidden_size).cuda()
+    opt = torch.optim.AdamW(adapter.parameters(), lr=LR, weight_decay=0.01)
+
+    train = [it for it in data if it["split"] == "train"
+             and it["label"] != "drop" and it["sol_code"]]
+    train = [it for it in train
+             if len(tok(mbpp_prompt(tok, it, DIRECT_SUFFIX))["input_ids"])
+             + len(tok(it["sol_code"])["input_ids"]) + 12 <= MAX_TOK]
+    val = train[:24]
+    pool = train[24:]
+    print(f"pool={len(pool)} val={len(val)}  lam={ARGS.lam} kmax={ARGS.kmax}",
+          flush=True)
+
+    log = []
+    t0 = time.time()
+    for step in range(STEPS):
+        batch = rng.sample(pool, BATCH)
+        ids, msk, lab, lmask = build_code_batch(tok, batch)
+        # last prompt position = last position where label == -100 & mask==1
+        pl = (lmask.long().cumsum(-1).argmax(-1))
+        for g in opt.param_groups:
+            g["lr"] = lr_at(step)
+        logits, q, ei = halting_loop_logits(
+            looper, adapter, ids, ARGS.kmax, msk, pl, loop_mask=lmask)
+        ce = F.cross_entropy(logits[:, :-1].flatten(0, 1).float(),
+                             lab[:, 1:].flatten(), ignore_index=-100)
+        loss = ce + lam_at(step) * ei.mean()
+        opt.zero_grad(set_to_none=True)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(adapter.parameters(), 1.0)
+        opt.step()
+        hard_b = [it["label"] == "hard" for it in batch]
+        log.append({"step": step, "ce": ce.item(), "ei": ei.mean().item(),
+                    "ei_hard": (ei[torch.tensor(hard_b)].mean().item()
+                                if any(hard_b) else None),
+                    "ei_easy": (ei[~torch.tensor(hard_b)].mean().item()
+                                if not all(hard_b) else None)})
+        if step % 10 == 0:
+            r = log[-1]
+            print(f"step {step:4d}  ce={r['ce']:.4f}  E[k]={r['ei']:.2f}  "
+                  f"(hard {r['ei_hard']} / easy {r['ei_easy']})  "
+                  f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True)
+        if step % 100 == 99 or step == STEPS - 1:
+            torch.save(adapter.state_dict(),
+                       OUT / f"adapter_{TAG}_e{step+1}.pt")
+            json.dump(log, open(OUT / f"train_{TAG}_log.json", "w"))
+    print("done", flush=True)
+
+
+if __name__ == "__main__":
+    main()