item 30 pre-registered (Nils's design): carried state as metacognitive signal — answer-readiness head on L30 at line boundaries, carry-vs-FF AUC as primary, theta sweep + oracle stop via item-20 LUT methodology
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Nils
Claude Fable 5
@@ -0,0 +1,12 @@
|
||||
# gpuq-in: results-loop/star_data.json results-loop/adapter_carrycot_e400.pt results-loop/adapter_carrycot_ff_e400.pt
|
||||
# gpuq-out: results-loop/metacog_*.pt results-loop/metacog_*.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 metacog.py harvest --split train
|
||||
$P metacog.py harvest --split test
|
||||
$P metacog.py harvest --split train --feedforward
|
||||
$P metacog.py harvest --split test --feedforward
|
||||
$P metacog.py head
|
||||
$P metacog.py sweep
|
||||
@@ -0,0 +1,245 @@
|
||||
"""Item 30: the carried state as a metacognitive signal (Nils's design).
|
||||
|
||||
The program's evidence: the carry's inference-time value is deferral /
|
||||
completion-state (430 divergence; drop-bucket p=0.0094), and gate quality
|
||||
is worth ~9.6 points with classifier quality binding (item 20). This
|
||||
exploits both: a small answer-readiness head g(s) on the carried L30
|
||||
state at scratchpad line boundaries. The frozen model reads nothing new
|
||||
— g's decisions reach it as tokens.
|
||||
|
||||
Stages:
|
||||
harvest --split train|test [--feedforward]
|
||||
Roll out rung-A generation (k=2, p by label), teacher-forced replay
|
||||
to capture carried states at line boundaries, then fork a forced
|
||||
"Answer: " completion per boundary and label it by correctness.
|
||||
Saves metacog_<split>[_ff].pt: states (N,1536), labels, item idx,
|
||||
boundary index, natural-end correctness.
|
||||
head
|
||||
Train class-balanced logistic + 2-layer MLP on train states; report
|
||||
test AUC for carry and (if present) FF states.
|
||||
sweep
|
||||
Offline theta sweep on the test LUT: answer at first boundary with
|
||||
g >= theta, else natural end. Accuracy / mean boundary / oracle.
|
||||
"""
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from carry_common import PAUSE_ID, _pos_ids, carry_logits, generate_carry_c
|
||||
from loop_common import (BandLooper, MergeAdapter, chat_prompt,
|
||||
DIRECT_SUFFIX, last_number, num_eq)
|
||||
|
||||
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"))
|
||||
K = 2
|
||||
P_BY_LABEL = {"easy": 2, "hard": 6, "drop": 6}
|
||||
NL_MAX_BOUND = 6
|
||||
|
||||
|
||||
def newline_positions(tok, gen_ids, start):
|
||||
"""Positions (absolute) of '\n' tokens within the generated span."""
|
||||
out = []
|
||||
for j, t in enumerate(gen_ids):
|
||||
if "\n" in tok.decode([t]):
|
||||
out.append(start + j)
|
||||
return out
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def harvest(args, model, tok, looper, adapter):
|
||||
items = [it for it in json.load(open(OUT / "star_data.json"))
|
||||
if it["split"] == args.split]
|
||||
if args.split == "test":
|
||||
items = items[:256]
|
||||
if args.limit:
|
||||
items = items[:args.limit]
|
||||
rows = []
|
||||
t0 = time.time()
|
||||
for bi in range(0, len(items), args.batch):
|
||||
chunk = items[bi:bi + args.batch]
|
||||
p = max(P_BY_LABEL[it["label"]] for it in chunk)
|
||||
enc = tok([chat_prompt(tok, it["question"], DIRECT_SUFFIX)
|
||||
for it in chunk], return_tensors="pt", padding=True,
|
||||
add_special_tokens=False).to("cuda")
|
||||
gen = generate_carry_c(looper, adapter, tok, enc["input_ids"],
|
||||
enc["attention_mask"], K, p,
|
||||
max_new_tokens=96,
|
||||
feedforward=args.feedforward)
|
||||
n_pre = enc["input_ids"].shape[1] + p
|
||||
# teacher-forced replay of each rollout to capture carried states
|
||||
for b, it in enumerate(chunk):
|
||||
seq = gen[b]
|
||||
# strip left padding
|
||||
keep = seq != (tok.pad_token_id or 0)
|
||||
keep[enc["attention_mask"].shape[1]:] = True
|
||||
seq = seq[keep]
|
||||
n_prompt = int(enc["attention_mask"][b].sum())
|
||||
# trim at end-of-turn
|
||||
gen_span = seq[n_prompt + p:]
|
||||
eot = tok.convert_tokens_to_ids("<end_of_turn>")
|
||||
ends = (gen_span == eot).nonzero()
|
||||
g_end = int(ends[0, 0]) if ends.numel() else len(gen_span)
|
||||
gen_span = gen_span[:g_end]
|
||||
full = seq[: n_prompt + p + g_end][None].cuda()
|
||||
msk = torch.ones_like(full)
|
||||
_, S = carry_logits(looper, adapter, full, msk,
|
||||
torch.tensor([n_prompt], device="cuda"), K,
|
||||
feedforward=args.feedforward,
|
||||
return_states=True)
|
||||
text = tok.decode(gen_span, skip_special_tokens=True)
|
||||
nat_ok = bool(num_eq(last_number(text), it["gold"]))
|
||||
nls = newline_positions(tok, gen_span.tolist(), n_prompt + p)
|
||||
nls = [x for x in nls if x < full.shape[1] - 2][:NL_MAX_BOUND]
|
||||
for k_b, pos in enumerate(nls):
|
||||
# forced answer from this prefix
|
||||
pre = seq[: pos + 1].cuda()
|
||||
fa = tok("Answer: ", add_special_tokens=False,
|
||||
return_tensors="pt")["input_ids"][0].cuda()
|
||||
inp = torch.cat([pre, fa])[None].cuda()
|
||||
am = torch.ones_like(inp)
|
||||
g2 = generate_carry_c(looper, adapter, tok, inp, am, 0, 0,
|
||||
max_new_tokens=8,
|
||||
feedforward=args.feedforward)
|
||||
ans = tok.decode(g2[0, inp.shape[1]:],
|
||||
skip_special_tokens=True)
|
||||
ok = bool(num_eq(last_number("Answer: " + ans), it["gold"]))
|
||||
rows.append({"idx": it["idx"], "label": it["label"],
|
||||
"bound": k_b, "ok": ok, "nat_ok": nat_ok,
|
||||
"state": S[0, pos].float().cpu()})
|
||||
if bi % (args.batch * 8) == 0:
|
||||
print(f"[{bi+len(chunk)}/{len(items)}] rows={len(rows)} "
|
||||
f"({time.time()-t0:.0f}s)", flush=True)
|
||||
states = torch.stack([r.pop("state") for r in rows])
|
||||
tag = f"metacog_{args.split}" + ("_ff" if args.feedforward else "")
|
||||
torch.save({"states": states, "rows": rows}, OUT / f"{tag}.pt")
|
||||
pos_rate = sum(r["ok"] for r in rows) / max(1, len(rows))
|
||||
print(f"saved {len(rows)} boundary rows -> {tag}.pt "
|
||||
f"(answer-ready rate {pos_rate:.3f})", flush=True)
|
||||
|
||||
|
||||
def auc(scores, labels):
|
||||
order = sorted(range(len(scores)), key=lambda i: scores[i])
|
||||
r = {}
|
||||
for rank, i in enumerate(order):
|
||||
r[i] = rank + 1
|
||||
pos = [i for i, l in enumerate(labels) if l]
|
||||
neg = [i for i, l in enumerate(labels) if not l]
|
||||
if not pos or not neg:
|
||||
return float("nan")
|
||||
s = sum(r[i] for i in pos)
|
||||
return (s - len(pos) * (len(pos) + 1) / 2) / (len(pos) * len(neg))
|
||||
|
||||
|
||||
def train_head(tr_s, tr_y, dev="cuda", hidden=64, steps=400, seed=0):
|
||||
torch.manual_seed(seed)
|
||||
d = tr_s.shape[1]
|
||||
net = torch.nn.Sequential(
|
||||
torch.nn.LayerNorm(d), torch.nn.Linear(d, hidden),
|
||||
torch.nn.GELU(), torch.nn.Linear(hidden, 1)).to(dev)
|
||||
opt = torch.optim.AdamW(net.parameters(), lr=1e-3, weight_decay=0.01)
|
||||
y = tr_y.float().to(dev)
|
||||
w = torch.where(y > 0, (1 - y.mean()) / y.mean().clamp(min=1e-3),
|
||||
torch.ones_like(y))
|
||||
x = tr_s.to(dev)
|
||||
for i in range(steps):
|
||||
logit = net(x).squeeze(-1)
|
||||
loss = F.binary_cross_entropy_with_logits(logit, y, weight=w)
|
||||
opt.zero_grad(); loss.backward(); opt.step()
|
||||
return net
|
||||
|
||||
|
||||
def head(args):
|
||||
def load(tag):
|
||||
d = torch.load(OUT / f"{tag}.pt", map_location="cpu")
|
||||
y = torch.tensor([r["ok"] for r in d["rows"]])
|
||||
return d["states"], y, d["rows"]
|
||||
results = {}
|
||||
for name, tr_tag, te_tag in [
|
||||
("carry", "metacog_train", "metacog_test"),
|
||||
("ff", "metacog_train_ff", "metacog_test_ff")]:
|
||||
if not (OUT / f"{tr_tag}.pt").exists():
|
||||
continue
|
||||
tr_s, tr_y, _ = load(tr_tag)
|
||||
te_s, te_y, te_rows = load(te_tag)
|
||||
net = train_head(tr_s, tr_y)
|
||||
with torch.no_grad():
|
||||
sc = torch.sigmoid(net(te_s.cuda()).squeeze(-1)).cpu()
|
||||
a = auc(sc.tolist(), te_y.tolist())
|
||||
results[name] = a
|
||||
print(f"{name}: test AUC={a:.3f} (n={len(te_y)}, "
|
||||
f"pos-rate={te_y.float().mean():.3f})", flush=True)
|
||||
if name == "carry":
|
||||
torch.save({"net": net.state_dict(),
|
||||
"scores": sc, }, OUT / "metacog_head.pt")
|
||||
for r, s_ in zip(te_rows, sc.tolist()):
|
||||
r["g"] = s_
|
||||
json.dump(te_rows, open(OUT / "metacog_lut.json", "w"))
|
||||
json.dump(results, open(OUT / "metacog_auc.json", "w"))
|
||||
|
||||
|
||||
def sweep(args):
|
||||
rows = json.load(open(OUT / "metacog_lut.json"))
|
||||
by_item = {}
|
||||
for r in rows:
|
||||
by_item.setdefault(r["idx"], []).append(r)
|
||||
n = len(by_item)
|
||||
nat = sum(1 for rs in by_item.values() if rs[0]["nat_ok"]) / n
|
||||
print(f"items={n} natural-end accuracy={nat:.3f}")
|
||||
print("theta acc mean-bound answered-early")
|
||||
for th in (0.5, 0.6, 0.7, 0.8, 0.85, 0.9, 0.95, 0.99):
|
||||
hits, bsum, early = 0, 0, 0
|
||||
for rs in by_item.values():
|
||||
rs = sorted(rs, key=lambda r: r["bound"])
|
||||
fired = next((r for r in rs if r["g"] >= th), None)
|
||||
if fired is not None:
|
||||
hits += fired["ok"]; bsum += fired["bound"] + 1; early += 1
|
||||
else:
|
||||
hits += rs[0]["nat_ok"]; bsum += len(rs) + 1
|
||||
print(f"{th:.2f} {hits/n:.3f} {bsum/n:.2f} {early}/{n}")
|
||||
oracle, obsum = 0, 0
|
||||
for rs in by_item.values():
|
||||
best = max([r["ok"] for r in rs] + [rs[0]["nat_ok"]])
|
||||
oracle += best
|
||||
obsum += min([r["bound"] + 1 for r in rs if r["ok"]],
|
||||
default=len(rs) + 1)
|
||||
print(f"ORACLE stop: acc={oracle/n:.3f} mean-bound={obsum/n:.2f}")
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("stage", choices=("harvest", "head", "sweep"))
|
||||
ap.add_argument("--split", default="train")
|
||||
ap.add_argument("--feedforward", action="store_true")
|
||||
ap.add_argument("--adapter", default=None)
|
||||
ap.add_argument("--batch", type=int, default=8)
|
||||
ap.add_argument("--limit", type=int, default=0)
|
||||
args = ap.parse_args()
|
||||
if args.stage == "harvest":
|
||||
model, tok = load_model(dtype=torch.bfloat16)
|
||||
tok.padding_side = "left"
|
||||
looper = BandLooper(model)
|
||||
adapter = MergeAdapter(
|
||||
d=model.config.get_text_config().hidden_size).cuda()
|
||||
default = OUT / ("adapter_carrycot_ff_e400.pt" if args.feedforward
|
||||
else "adapter_carrycot_e400.pt")
|
||||
adapter.load_state_dict(torch.load(args.adapter or default,
|
||||
map_location="cuda"))
|
||||
adapter.eval()
|
||||
harvest(args, model, tok, looper, adapter)
|
||||
elif args.stage == "head":
|
||||
head(args)
|
||||
else:
|
||||
sweep(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user