246 lines
10 KiB
Python
246 lines
10 KiB
Python
"""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()
|