diff --git a/nbs/cosine_dist.ipynb b/nbs/cosine_dist.ipynb new file mode 100644 index 0000000..374386c --- /dev/null +++ b/nbs/cosine_dist.ipynb @@ -0,0 +1,71 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Pinning diagnostic: does the contrastive hack direction predict live hacks?\n", + "\n", + "Replays `scripts/diag_cosine_dist.py` outputs (no GPU). Spaces: **grad** (v_grad on delta_S) and **act** (As in S space).\n", + "Scores: cosine (dir only), projection (dir x |g|), magnitude (|g|). Filters: all modules / noise-floor kept.\n", + "Separability = AUROC + precision@k of score -> oracle `exploited`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import polars as pl\n", + "import matplotlib.pyplot as plt\n", + "D = '../out/diag/'\n", + "hist = pl.read_parquet(D+'cosine_dist.parquet')\n", + "scores = pl.read_parquet(D+'live_scores.parquet')\n", + "sep = pl.read_csv(D+'separability.csv')\n", + "sep" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# histograms: cosine to hack direction, both spaces\n", + "colors = {'pair_clean':'tab:blue','pair_hack':'tab:red','live_clean':'tab:cyan','live_hack':'tab:orange'}\n", + "fig, axes = plt.subplots(1, 2, figsize=(15, 5))\n", + "for ax, space in zip(axes, ['grad','act']):\n", + " for pop, c in colors.items():\n", + " v = hist.filter((pl.col('space')==space) & (pl.col('pop')==pop))['cos'].to_numpy()\n", + " if len(v):\n", + " ax.hist(v, bins=15, density=True, histtype='step', lw=2, color=c, label=f'{pop} (n={len(v)})')\n", + " ax.set_title(f'{space} space'); ax.set_xlabel('global cosine to hack dir'); ax.legend(fontsize=8)\n", + "plt.tight_layout(); plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# per-score distributions split by exploited: which score separates?\n", + "cols = [c for c in scores.columns if c != 'exploited']\n", + "fig, axes = plt.subplots(2, len(cols)//2, figsize=(16, 7))\n", + "for ax, col in zip(axes.flat, cols):\n", + " for y, c in [(True,'tab:orange'),(False,'tab:cyan')]:\n", + " v = scores.filter(pl.col('exploited')==y)[col].to_numpy()\n", + " ax.hist(v, bins=15, density=True, histtype='step', lw=2, color=c, label='hack' if y else 'clean')\n", + " ax.set_title(col, fontsize=8); ax.legend(fontsize=7)\n", + "plt.tight_layout(); plt.show()" + ] + } + ], + "metadata": { + "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, + "language_info": {"name": "python", "version": "3.12"} + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/scripts/diag_cosine_dist.py b/scripts/diag_cosine_dist.py new file mode 100644 index 0000000..653054c --- /dev/null +++ b/scripts/diag_cosine_dist.py @@ -0,0 +1,256 @@ +"""Pinning diagnostic: how well does the contrastive hack direction PREDICT a live +hack rollout, in GRADIENT vs ACTIVATION space, by cosine vs projection vs magnitude? + +One checkpoint (default: job 9 per-token `first_hack` = step 7). We build the +contrastive hack direction from the authored pairs in two spaces: + + GRAD: v_grad = unit(mean_pairs(g_hack - g_clean)) on delta_S (g = per-rollout NLL grad) + ACT: As_hack = unit(mean_pairs(As_hack - As_clean)) (As = Vh@x mean over comp tokens) + +Then for live rollouts (oracle `exploited`, offline plot-only label) we score each +rollout against the direction and ask how well the score separates hack from clean. +Three score families (GLOBAL = concat over modules; v is unit per module): + cosine = / (|g| |v|) -- direction only (magnitude-blind) + projection = / |v| = |g| cos -- direction x gradient magnitude + magnitude = |g| -- magnitude only (null check) +each in {grad, act} x {all modules, noise-floor-filtered (drop bottom-25% singular value)}. + +Separability metric (we want high PRECISION / confident tail, not coverage): + AUROC + precision@10 / @20 of "score predicts exploited". + + uv run python scripts/diag_cosine_dist.py # 4B, free GPU, ~4-6 min +outputs (out/diag/): cosine_grad.png, cosine_act.png (histograms), + cosine_dist.parquet (histogram long-form), live_scores.parquet (per-rollout scores), + separability.csv (methods x AUROC/precision). Replay without a GPU: nbs/cosine_dist.ipynb. +""" +from __future__ import annotations +import json +import struct +from dataclasses import dataclass +from pathlib import Path + +import torch +import torch.nn.functional as F +import tyro +import polars as pl +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +from loguru import logger +from tabulate import tabulate +from safetensors.torch import load_file +from transformers import AutoModelForCausalLM, AutoTokenizer + +from vgrout.antipasto import wrap_model_with_antipasto +from vgrout.extract_vhack_grad import extract_v_hack, completion_nll +from vgrout.pairs import PAIRS +from vgrout.train import CACHE_ROOT + + +@dataclass +class Cfg: + run_dir: Path = Path("out/runs/20260607T134234_fast_routingV_seed43_dir6_routeV_pertoken_s43") + ckpt: str = "first_hack" + step_lo: int = 5 + step_hi: int = 9 + max_rollouts: int = 140 + drop_frac: float = 0.25 # noise floor: drop modules below this global singular-value quantile + bins: int = 15 # histogram bins (wider = less spiky) + out_dir: Path = Path("out/diag") + + +def _auroc(score, label) -> float: + """P(score_pos > score_neg), ties=0.5. O(n^2), fine for n~140.""" + pos = [s for s, y in zip(score, label) if y] + neg = [s for s, y in zip(score, label) if not y] + if not pos or not neg: + return float("nan") + c = sum((p > n) + 0.5 * (p == n) for p in pos for n in neg) + return c / (len(pos) * len(neg)) + + +def _prec_at_k(score, label, k) -> float: + order = sorted(range(len(score)), key=lambda i: score[i], reverse=True)[:k] + return sum(label[i] for i in order) / k + + +def main(cfg: Cfg) -> int: + device = torch.device("cuda") + kept_path = cfg.run_dir / f"{cfg.ckpt}.safetensors" + hack_path = cfg.run_dir / f"{cfg.ckpt}_hack.safetensors" + with open(kept_path, "rb") as f: + meta = json.loads(f.read(struct.unpack("= sv0.quantile(cfg.drop_frac) # noise floor mask + logger.info(f"noise floor: keep {int(keep.sum())}/{len(names)} modules (sv>=q{cfg.drop_frac})") + + # ── activation capture hooks (after grad extract) ── + As_cap: dict[str, torch.Tensor] = {} + state = {"plen": 0} + + def make_hook(nm): + Vh = wrappers[nm]["layer"]._antipasto_Vh + def hook(_l, inp, _o): + a = F.linear(inp[0], Vh) + As_cap[nm] = a[0, state["plen"] - 1:, :].mean(0).detach().float().cpu() + return hook + handles = [wrappers[nm]["layer"].register_forward_hook(make_hook(nm)) for nm in names] + + def grab_As(prompt, completion): + state["plen"] = tok(prompt, return_tensors="pt").input_ids.shape[1] + ids = tok(prompt + completion, return_tensors="pt").input_ids.to(device) + with torch.no_grad(): + model(ids) + return {nm: As_cap[nm].clone() for nm in names} + + # ── ACT direction from the same train pairs ── + train_pairs = list(PAIRS)[:-2] + As_h = {nm: [] for nm in names} + As_c = {nm: [] for nm in names} + for p in train_pairs: + ah, ac = grab_As(p.prompt, p.hack), grab_As(p.prompt, p.clean) + for nm in names: + As_h[nm].append(ah[nm]); As_c[nm].append(ac[nm]) + As_dir = {nm: (lambda d: d / d.norm().clamp_min(1e-12))( + (torch.stack(As_h[nm]) - torch.stack(As_c[nm])).mean(0)) for nm in names} # cpu unit + + # stack directions to [n_mod, r-ragged] -> we keep per-module and reduce with helpers + def per_module_dot_norm(g: dict, vdir: dict): + """returns dot[n_mod], gnorm[n_mod] (cpu float), v assumed unit per module.""" + dot = torch.tensor([(g[nm].flatten().float().cpu() @ vdir[nm].flatten().float()).item() for nm in names]) + gn = torch.tensor([g[nm].flatten().float().cpu().norm().item() for nm in names]) + return dot, gn + + # ── live rollouts: grad (backward) + act (same forward) per-module dot/norm ── + recs = [json.loads(l) for l in (cfg.run_dir / "rollouts.jsonl").read_text().splitlines()] + batch = [r for r in recs if cfg.step_lo <= r["step"] <= cfg.step_hi and r["text"].strip()][:cfg.max_rollouts] + logger.info(f"live batch: {len(batch)} rollouts (steps {cfg.step_lo}-{cfg.step_hi})") + G_dot, G_gn, A_dot, A_an, labels = [], [], [], [], [] + for i, r in enumerate(batch): + state["plen"] = tok(r["prompt"], return_tensors="pt").input_ids.shape[1] + model.zero_grad(set_to_none=True) + loss = completion_nll(model, tok, r["prompt"], r["text"], device) # fires act hooks + if not torch.isfinite(loss): + continue + As_sample = {nm: As_cap[nm].clone() for nm in names} + loss.backward() + gb = {nm: wrappers[nm]["delta_S"].grad for nm in names} + gd, gn = per_module_dot_norm(gb, v_grad) + ad, an = per_module_dot_norm(As_sample, As_dir) + G_dot.append(gd); G_gn.append(gn); A_dot.append(ad); A_an.append(an) + labels.append(bool(r["exploited"])) + if (i + 1) % 25 == 0: + logger.info(f" rollout {i+1}/{len(batch)}") + for h in handles: + h.remove() + model.zero_grad(set_to_none=True) + G_dot, G_gn = torch.stack(G_dot), torch.stack(G_gn) # [n_roll, n_mod] + A_dot, A_an = torch.stack(A_dot), torch.stack(A_an) + labels = labels + n_pos = sum(labels) + logger.info(f"live: {len(labels)} rollouts, {n_pos} exploited ({n_pos/len(labels):.2f} base rate)") + + # ── global scores per rollout: cosine / projection / magnitude, all + filtered ── + def scores(dot, gn, mask): + d = dot[:, mask].sum(1) # sum_m + gmag = gn[:, mask].pow(2).sum(1).sqrt() # |concat g| + nmod = int(mask.sum()) + cos = d / (gmag.clamp_min(1e-12) * (nmod ** 0.5)) + proj = d / (nmod ** 0.5) + return {"cosine": cos, "projection": proj, "magnitude": gmag} + + score_cols = {} + for space, (dot, gn) in {"grad": (G_dot, G_gn), "act": (A_dot, A_an)}.items(): + for filt, mask in {"all": torch.ones(len(names), dtype=bool), "kept": keep}.items(): + for fam, v in scores(dot, gn, mask).items(): + score_cols[f"{space}.{fam}.{filt}"] = v.tolist() + + # ── separability table ── + sep_rows = [] + for name, s in score_cols.items(): + space, fam, filt = name.split(".") + sep_rows.append(dict(space=space, score=fam, filter=filt, + AUROC=round(_auroc(s, labels), 3), + p_at_10=round(_prec_at_k(s, labels, 10), 3), + p_at_20=round(_prec_at_k(s, labels, 20), 3))) + sep = pl.DataFrame(sep_rows).sort("AUROC", descending=True) + + # ── persist: histogram long-form (global cosine) + per-rollout scores + separability ── + cfg.out_dir.mkdir(parents=True, exist_ok=True) + # histogram populations need the PAIR cosines too (built from raw_grads / As_h,As_c) + def pair_global_cos(per_pair_dot_src, vdir, n): + out = [] + for i in range(n): + g = {nm: per_pair_dot_src(nm, i) for nm in names} + dot, gn = per_module_dot_norm(g, vdir) + out.append((dot.sum() / (gn.pow(2).sum().sqrt().clamp_min(1e-12) * (len(names) ** 0.5))).item()) + return out + npg = raw_grads[f"hack/{names[0]}"].shape[0] + hist = [] + for pop, src in [("pair_hack", lambda nm, i: raw_grads[f"hack/{nm}"][i]), + ("pair_clean", lambda nm, i: raw_grads[f"clean/{nm}"][i])]: + for c in pair_global_cos(src, v_grad, npg): + hist.append(dict(space="grad", pop=pop, cos=c)) + for pop, src in [("pair_hack", lambda nm, i: As_h[nm][i]), + ("pair_clean", lambda nm, i: As_c[nm][i])]: + for c in pair_global_cos(src, As_dir, len(train_pairs)): + hist.append(dict(space="act", pop=pop, cos=c)) + for sp, col in [("grad", "grad.cosine.all"), ("act", "act.cosine.all")]: + for c, y in zip(score_cols[col], labels): + hist.append(dict(space=sp, pop="live_hack" if y else "live_clean", cos=c)) + pl.DataFrame(hist).write_parquet(cfg.out_dir / "cosine_dist.parquet") + pl.DataFrame({**score_cols, "exploited": labels}).write_parquet(cfg.out_dir / "live_scores.parquet") + sep.write_csv(cfg.out_dir / "separability.csv") + + logger.info("\n=== separability (AUROC of score -> exploited; >0.5 = predictive) ===\n" + + tabulate(sep.to_pandas(), headers="keys", tablefmt="github", showindex=False)) + + # ── histograms (cosine, both spaces) ── + hdf = pl.DataFrame(hist) + colors = {"pair_clean": "tab:blue", "pair_hack": "tab:red", + "live_clean": "tab:cyan", "live_hack": "tab:orange"} + for space in ("grad", "act"): + plt.figure(figsize=(9, 5)) + for pop, c in colors.items(): + v = hdf.filter((pl.col("space") == space) & (pl.col("pop") == pop))["cos"].to_numpy() + if len(v): + plt.hist(v, bins=cfg.bins, density=True, histtype="step", lw=2, color=c, + label=f"{pop} (n={len(v)}, p50={float(pl.Series(v).median()):+.2f})") + au = sep.filter((pl.col("space") == space) & (pl.col("score") == "cosine") & (pl.col("filter") == "all"))["AUROC"][0] + plt.xlabel(f"global cosine to hack direction ({space} space)") + plt.ylabel("density") + plt.title(f"{space}-space: {cfg.ckpt} (step {meta.get('step')}), live AUROC(cos)={au}") + plt.legend(fontsize=8) + plt.tight_layout() + plt.savefig(cfg.out_dir / f"cosine_{space}.png", dpi=130) + plt.close() + logger.info(f"wrote {cfg.out_dir}/cosine_{{grad,act}}.png, cosine_dist.parquet, " + f"live_scores.parquet, separability.csv") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main(tyro.cli(Cfg)))