item 20: gate threshold curve + oracle bound (probs pass, merge per-item LUT, offline sweep)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -439,3 +439,18 @@ and compute savings are demonstrated and cheap; difficulty-selective
|
|||||||
DEPTH allocation remains unsolved at 3K-param-head scale — binding
|
DEPTH allocation remains unsolved at 3K-param-head scale — binding
|
||||||
constraint is classifier quality on the k=0/s0 state, exactly where E0
|
constraint is classifier quality on the k=0/s0 state, exactly where E0
|
||||||
started. Next per plan: threshold sweep (cheap) before any E2.
|
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.
|
||||||
|
|||||||
@@ -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()
|
||||||
@@ -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")
|
||||||
@@ -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
|
||||||
Reference in New Issue
Block a user