From 4b9e168838ed75879289ba0c8617a28ff63cb079 Mon Sep 17 00:00:00 2001 From: Nils Date: Thu, 16 Jul 2026 02:03:12 +0200 Subject: [PATCH] item 20: gate threshold curve + oracle bound (probs pass, merge per-item LUT, offline sweep) Co-Authored-By: Claude Fable 5 --- results-loop/PROTOCOL_UNIFIED.md | 15 +++++++ scripts/gate_probe_probs.py | 76 ++++++++++++++++++++++++++++++++ scripts/gate_threshold_curve.py | 74 +++++++++++++++++++++++++++++++ scripts/jobs/zzz_k_gatecurve.sh | 9 ++++ 4 files changed, 174 insertions(+) create mode 100644 scripts/gate_probe_probs.py create mode 100644 scripts/gate_threshold_curve.py create mode 100644 scripts/jobs/zzz_k_gatecurve.sh diff --git a/results-loop/PROTOCOL_UNIFIED.md b/results-loop/PROTOCOL_UNIFIED.md index ce3bd26..de66385 100644 --- a/results-loop/PROTOCOL_UNIFIED.md +++ b/results-loop/PROTOCOL_UNIFIED.md @@ -439,3 +439,18 @@ and compute savings are demonstrated and cheap; difficulty-selective DEPTH allocation remains unsolved at 3K-param-head scale — binding constraint is classifier quality on the k=0/s0 state, exactly where E0 started. Next per plan: threshold sweep (cheap) before any E2. + +20. **E1 threshold curve + oracle bound (pre-registered 2026-07-16 ~02:15, + before running).** Phase 1: record E1c head's halt probabilities per + test item (one GPU pass). Phase 2: per-item outcomes for the frozen + curriculum merge at k=0/1/2/4 (four generation sweeps, tag merge_lut — + doubles as the reusable gate-evaluation lookup table and supplies the + long-missing per-item logs for the canonical merge). Phase 3 (offline): + gated accuracy at thresholds .3-.99 by composing k*(theta) with the + lookup; plus the ORACLE gate (best k per item) = the ceiling any gate + can reach with this merge. Predictions: (a) some theta gives hard >= + 32% with easy >= 93% and E[k] <= 2.2 (dominating E0 on compute at + comparable accuracy); (b) the curve is monotone in theta; (c) oracle + overall >= 55% — if so, gate-quality headroom is large and further + gate work is justified; if oracle < 53%, gating this merge is nearly + saturated and the program pivots to E2 or closes. diff --git a/scripts/gate_probe_probs.py b/scripts/gate_probe_probs.py new file mode 100644 index 0000000..a30e21d --- /dev/null +++ b/scripts/gate_probe_probs.py @@ -0,0 +1,76 @@ +"""Item 20 phase 1: record the E1c halting head's probabilities per item. + +One pass over the test set with the gatehead adapter: p0 (pre-loop, on +s_0) and p_1..p_kmax (after each iteration). With these, k*(threshold) +is computable offline for any threshold — no more GPU passes. +""" + +import argparse +import json +import os +import sys +from pathlib import Path + +import torch + +from halting_common import HaltingMergeAdapter +from loop_common import BandLooper +from prep_mbpp import DIRECT_SUFFIX, mbpp_prompt + +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("--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] + rows = [] + 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") + ids, msk = enc["input_ids"], enc["attention_mask"] + calls, _ = looper.capture(ids, msk, logits_to_keep=1) + e = looper._hin[looper.l0].detach() + B = e.shape[0] + bidx = torch.arange(B, device=e.device) + pl = torch.full((B,), ids.shape[1] - 1, device=e.device, + dtype=torch.long) + s = looper.band(e, calls) + probs = [adapter.halt_prob(e[bidx, pl], s[bidx, pl])] + for _ in range(args.kmax): + x = adapter(e, s) + s = looper.band(x, calls) + probs.append(adapter.halt_prob(e[bidx, pl], s[bidx, pl])) + P = torch.stack(probs, -1).float().cpu() # (B, kmax+1) + for j, it in enumerate(chunk): + rows.append({"task_id": it["task_id"], "label": it["label"], + "p": [round(v, 5) for v in P[j].tolist()]}) + if i % 40 == 0: + print(f"[{i+len(chunk)}/{len(items)}]", flush=True) + json.dump(rows, open(OUT / "gate_probs.json", "w"), indent=1) + print("wrote", OUT / "gate_probs.json") + + +if __name__ == "__main__": + main() diff --git a/scripts/gate_threshold_curve.py b/scripts/gate_threshold_curve.py new file mode 100644 index 0000000..659b171 --- /dev/null +++ b/scripts/gate_threshold_curve.py @@ -0,0 +1,74 @@ +"""Item 20 phase 3 (offline, no GPU): threshold sweep + oracle bound. + +Inputs: gate_probs.json (per-item halt probabilities, phase 1) and +eval_code_merge_lut.json (per-item outcomes at k=0/1/2/4, phase 2). +For each threshold theta: k*(item) = first point where cumulative halt +mass >= theta (pre-loop consult included, remainder to kmax), outcome +looked up per item; reports overall/easy/hard accuracy and E[k]. +Also: the oracle gate (best k per item) — the gating ceiling. +""" + +import json +import os +from pathlib import Path + +OUT = Path(os.environ.get("LOOP_OUT", + Path(__file__).resolve().parent.parent / "results-loop")) + +probs = {r["task_id"]: r for r in json.load(open(OUT / "gate_probs.json"))} +lut = json.load(open(OUT / "eval_code_merge_lut.json")) +# lut["ks"][k]["per_item"] -> task_id, ok +ok = {} +for k, v in lut["ks"].items(): + for r in v["per_item"]: + ok.setdefault(r["task_id"], {})[int(k)] = r["ok"] +KS = sorted(int(k) for k in lut["ks"]) # e.g. [0, 1, 2, 4] +KMAX = max(KS) + + +def kstar(p, theta): + """p = [p0, p1, ..., pkmax]; deploy semantics of halted_k_per_item.""" + cum, keep = 0.0, 1.0 + for i, pi in enumerate(p): + cum += keep * pi + keep *= (1 - pi) + if cum >= theta: + return 0 if i == 0 else min(i, KMAX) + return KMAX + + +def nearest_k(k): + """Map k* to the nearest evaluated k (lut has only KS).""" + return min(KS, key=lambda x: (abs(x - k), x)) + + +def score(assign): + n, hit = {}, {} + ek = 0.0 + for tid, k in assign.items(): + lbl = probs[tid]["label"] + n[lbl] = n.get(lbl, 0) + 1 + hit[lbl] = hit.get(lbl, 0) + int(ok[tid][nearest_k(k)]) + ek += k + tot_n = sum(n.values()) + tot = sum(hit.values()) / tot_n + return tot, {l: hit[l] / n[l] for l in n}, ek / tot_n + + +print(f"{'theta':>6} {'overall':>8} {'easy':>6} {'hard':>6} {'E[k]':>5}") +best = [] +for theta in (0.3, 0.5, 0.7, 0.8, 0.9, 0.95, 0.98, 0.99): + assign = {tid: kstar(r["p"], theta) for tid, r in probs.items()} + tot, by, ek = score(assign) + print(f"{theta:>6} {tot:>8.3f} {by.get('easy',0):>6.3f} " + f"{by.get('hard',0):>6.3f} {ek:>5.2f}") + best.append({"theta": theta, "overall": tot, "by_label": by, "ek": ek}) + +oracle = {tid: max(KS, key=lambda k: (ok[tid][k], -k)) for tid in probs} +tot, by, ek = score(oracle) +print(f"{'oracle':>6} {tot:>8.3f} {by.get('easy',0):>6.3f} " + f"{by.get('hard',0):>6.3f} {ek:>5.2f}") +json.dump({"curve": best, "oracle": {"overall": tot, "by_label": by, + "ek": ek}}, + open(OUT / "gate_threshold_curve.json", "w"), indent=1) +print("wrote", OUT / "gate_threshold_curve.json") diff --git a/scripts/jobs/zzz_k_gatecurve.sh b/scripts/jobs/zzz_k_gatecurve.sh new file mode 100644 index 0000000..d2455d4 --- /dev/null +++ b/scripts/jobs/zzz_k_gatecurve.sh @@ -0,0 +1,9 @@ +# gpuq-in: results-loop/mbpp_data.json +# gpuq-out: results-loop/gate_probs.json results-loop/eval_code_merge_lut.json results-loop/gate_threshold_curve.json +git pull origin main -q 2>/dev/null +P=/home/nils/jspace/.venv/bin/python +export JLENS_MODEL=google/gemma-4-E2B-it LOOP_OUT=/home/nils/jspace/results-loop +cd /home/nils/jspace/scripts +$P gate_probe_probs.py --adapter $LOOP_OUT/adapter_gatehead_e300.pt --kmax 4 --n 250 +$P eval_loop_code.py --adapter $LOOP_OUT/adapter_code.pt --tag merge_lut --ks 0,1,2,4 --n 250 +$P gate_threshold_curve.py