| """Where does the model know that its current answer is right? Linear read-out per layer. |
| |
| At every probe cut of a --record-all shadow run (label: forced answer at that cut is correct), the residual stream of |
| every layer is recorded; for each layer an L2-regularised logistic regression is evaluated with trajectory-grouped |
| 5-fold cross-validation (no trajectory in both train and test). Also tests a temporal feature: the change of the |
| state since the previous cut (consistency of the internal state over time). |
| usage: layer_probe.py --rec results/rec/raw_clean.jsonl --items "data/coc/*.jsonl" --exclude data/coc/mini_math.jsonl |
| """ |
| import argparse, glob, json, os, sys |
| import torch |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| import transformers.integrations.sdpa_attention as _sdpa |
| _sdpa.use_gqa_in_sdpa = lambda *a, **k: False |
| from grade import grade_one |
|
|
| ap = argparse.ArgumentParser() |
| ap.add_argument("--rec", required=True) |
| ap.add_argument("--items", nargs="+", required=True) |
| ap.add_argument("--exclude", nargs="*", default=[]) |
| ap.add_argument("--model", default="models/LFM3-1.7B") |
| ap.add_argument("--cache", default="results/halt/layer_feats.pt") |
| args = ap.parse_args() |
|
|
| if not os.path.exists(args.cache): |
| tok = AutoTokenizer.from_pretrained(args.model) |
| model = AutoModelForCausalLM.from_pretrained(args.model, dtype=torch.bfloat16, attn_implementation="sdpa", |
| trust_remote_code=True).cuda().eval() |
| L = model.config.num_hidden_layers |
| cap = {} |
| for k, layer in enumerate(model.model.layers): |
| layer.register_forward_hook(lambda m, i, o, k=k: cap.__setitem__(k, (o[0] if isinstance(o, tuple) else o)[0])) |
| ends = {i for i in range(len(tok)) if (t := tok.convert_ids_to_tokens(i)) is not None |
| and (t.endswith(".") or t.endswith("Ċ"))} |
| items = {} |
| for pat in args.items: |
| for f in glob.glob(pat): |
| for l in open(f, encoding="utf-8"): |
| r = json.loads(l); items[r["id"]] = r |
| excl = set() |
| for f in args.exclude: |
| excl |= {json.loads(l)["id"] for l in open(f, encoding="utf-8")} |
| X, Y, T, TASK, IDX = [], [], [], [], [] |
| n = 0 |
| for l in open(args.rec, encoding="utf-8"): |
| r = json.loads(l) |
| if r["id"] in excl or r["id"] not in items or len(r.get("probe_log") or []) == 0: |
| continue |
| it = items[r["id"]] |
| p_ids = tok.apply_chat_template([{"role": "user", "content": it["prompt"]}], tokenize=True, |
| add_generation_prompt=True, return_dict=False) |
| resp = tok(r["completion"][len("<think>"):], add_special_tokens=False)["input_ids"][:8192] |
| pos, ys = [], [] |
| for p in r["probe_log"]: |
| c = min(p["cut"], len(resp)) |
| while c > 1 and resp[c - 1] not in ends: |
| c -= 1 |
| if c > 1: |
| pos.append(len(p_ids) + c - 1) |
| ys.append(float(grade_one({**r, "completion": "<think>x</think>" + p["text"], "n_tokens": 0}, it)["correct"])) |
| if not pos: |
| continue |
| with torch.no_grad(): |
| model(input_ids=torch.tensor([p_ids + resp], device="cuda"), use_cache=False) |
| X.append(torch.stack([cap[k][pos].float().cpu() for k in range(L)], 1)) |
| Y += ys; T += [n] * len(pos); TASK += [it["task"]] * len(pos); IDX += list(range(len(pos))) |
| n += 1 |
| torch.save({"X": torch.cat(X).half(), "y": torch.tensor(Y), "traj": torch.tensor(T), "task": TASK, |
| "idx": torch.tensor(IDX)}, args.cache) |
| d = torch.load(args.cache) |
| X, y, traj, idx = d["X"].float(), d["y"], d["traj"], d["idx"] |
| print(X.shape[0], "cut points,", int(y.sum()), "positive,", int(traj.max()) + 1, "trajectories") |
| prev = torch.zeros_like(X) |
| for i in range(1, len(X)): |
| if traj[i] == traj[i - 1]: |
| prev[i] = X[i - 1] |
| has_prev = idx > 0 |
|
|
|
|
| def auc(s, yy): |
| o = torch.argsort(s); r = torch.empty(len(s)); r[o] = torch.arange(1, len(s) + 1, dtype=torch.float) |
| n1 = yy.sum(); n0 = len(yy) - n1 |
| return ((r[yy == 1].sum() - n1 * (n1 + 1) / 2) / (n1 * n0)).item() |
|
|
|
|
| def cv_auc(F, yy, tr_ids, lam): |
| folds = torch.tensor([hash(int(t)) % 5 for t in tr_ids]) |
| scores = torch.zeros(len(yy)) |
| for f in range(5): |
| te = folds == f; trn = ~te |
| mu, sd = F[trn].mean(0), F[trn].std(0) + 1e-6 |
| A = (F - mu) / sd |
| w = torch.zeros(A.shape[1], requires_grad=True); b = torch.zeros(1, requires_grad=True) |
| opt = torch.optim.LBFGS([w, b], max_iter=200, line_search_fn="strong_wolfe") |
| def closure(): |
| opt.zero_grad() |
| loss = torch.nn.functional.binary_cross_entropy_with_logits(A[trn] @ w + b, yy[trn]) + lam * w.pow(2).sum() |
| loss.backward(); return loss |
| opt.step(closure) |
| scores[te] = (A[te] @ w + b).detach() |
| return auc(scores, yy) |
|
|
|
|
| print("layer AUC(state) AUC(state - previous cut) [5-fold, grouped by trajectory]") |
| for k in range(0, X.shape[1], 2): |
| a1 = max(cv_auc(X[:, k], y, traj, lam) for lam in (1e-1, 1.0)) |
| m = has_prev |
| a2 = max(cv_auc(torch.cat([X[m, k], X[m, k] - prev[m, k]], 1), y[m], traj[m], lam) for lam in (1e-1, 1.0)) |
| print(f"{k:5d} {a1:10.3f} {a2:10.3f}", flush=True) |
|
|
