File size: 5,345 Bytes
67ebcef | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | """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)) # [n_pos, L, H]
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)
|