diff --git a/scripts/diag_cosine_dist.py b/scripts/diag_cosine_dist.py index 02d7d10..2136e4a 100644 --- a/scripts/diag_cosine_dist.py +++ b/scripts/diag_cosine_dist.py @@ -73,6 +73,10 @@ def _prec_at_k(score, label, k) -> float: return sum(label[i] for i in order) / k +def _concat(d: dict, names: list) -> torch.Tensor: + return torch.cat([d[n].flatten().float().cpu() for n in names]) + + def main(cfg: Cfg) -> int: device = torch.device("cuda") kept_path = cfg.run_dir / f"{cfg.ckpt}.safetensors" @@ -102,15 +106,16 @@ def main(cfg: Cfg) -> int: top_k=1, tau_axis=0.0, n_heldout=2, device=device) v_grad = {nm: (lambda d: (d / d.norm().clamp_min(1e-12)))( (raw_grads[f"hack/{nm}"] - raw_grads[f"clean/{nm}"]).mean(0)) for nm in names} # cpu unit - sv0 = torch.tensor([v_sv[nm][0].item() for nm in names]) # [n_mod] - # filter levels by per-module contrastive singular value S0 (the noise floor signal): - # all = no filter (every module) - # keep75 = drop bottom 25% (training default) - # top25 = keep only the top 25% strongest-separating modules (strong filter) - masks = {"all": torch.ones(len(names), dtype=bool), - "keep75": sv0 >= sv0.quantile(0.25), - "top25": sv0 >= sv0.quantile(0.75)} - logger.info("filter levels: " + ", ".join(f"{k}={int(m.sum())}/{len(names)}" for k, m in masks.items())) + sv0 = torch.tensor([v_sv[nm][0].item() for nm in names]) # [n_mod] grad contrastive |D_m| (=S0) + + # filter levels keep the strongest-separating modules by a per-space weight w (=|D_m|): + # all/keep75/top25/top15/top05 = keep modules with w >= q{0, .25, .75, .85, .95}. + def make_masks(w): + return {"all": torch.ones(len(names), dtype=bool), + "keep75": w >= w.quantile(0.25), + "top25": w >= w.quantile(0.75), + "top15": w >= w.quantile(0.85), + "top05": w >= w.quantile(0.95)} # ── activation capture hooks (after grad extract) ── As_cap: dict[str, torch.Tensor] = {} @@ -139,8 +144,9 @@ def main(cfg: Cfg) -> int: 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 + As_D = {nm: (torch.stack(As_h[nm]) - torch.stack(As_c[nm])).mean(0) for nm in names} # act contrastive + As_dir = {nm: As_D[nm] / As_D[nm].norm().clamp_min(1e-12) for nm in names} # cpu unit + act_w = torch.tensor([As_D[nm].norm().item() for nm in names]) # act |D_m| # stack directions to [n_mod, r-ragged] -> we keep per-module and reduce with helpers def per_module_dot_norm(g: dict, vdir: dict): @@ -154,6 +160,7 @@ def main(cfg: Cfg) -> int: 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 = [], [], [], [], [] + Gvec, Avec = [], [] # concat per-rollout g / As (fp16) for the ideal-direction ceiling 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) @@ -166,6 +173,7 @@ def main(cfg: Cfg) -> int: 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) + Gvec.append(_concat(gb, names).half()); Avec.append(_concat(As_sample, names).half()) labels.append(bool(r["exploited"])) if (i + 1) % 25 == 0: logger.info(f" rollout {i+1}/{len(batch)}") @@ -178,19 +186,52 @@ def main(cfg: Cfg) -> int: 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| + # ── global scores per rollout: cosine / projection / magnitude / vote, per filter ── + # vote = module-weighted average of per-module cosines, weight w=|D_m| (soft version of + # the hard module filter: down-weight weak modules instead of cutting them). + def scores(dot, gn, mask, w): + dm, gm = dot[:, mask], gn[:, mask] + d = dm.sum(1) # sum_m + gmag = gm.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} + cos_m = dm / gm.clamp_min(1e-12) # per-module cosine [n_roll, n_kept] + wm = w[mask] + vote = (cos_m * wm).sum(1) / wm.sum().clamp_min(1e-12) + return {"cosine": cos, "projection": proj, "magnitude": gmag, "vote": vote} + # ── ceiling: ideal direction (mu_hack - mu_clean) fit on live rollouts, 2-fold CV ── + # The oracle-labelled best linear direction -- diagnostic upper bound only (never trains). + # If ceiling >> our pair-direction AUROC, the pairs are the problem (on-distribution pairs + # would help); if ceiling is also ~0.7, the distributions fundamentally overlap. + def ideal_auroc(vecs): + X = torch.stack(vecs).float() # [n, D] + y = torch.tensor(labels) + sc = torch.zeros(len(y)) + for fold in (0, 1): # 2-fold CV (fit on ~half, score other) + te = (torch.arange(len(y)) % 2 == fold) + tr = ~te + vd = X[tr][y[tr]].mean(0) - X[tr][~y[tr]].mean(0) + vd = vd / vd.norm().clamp_min(1e-12) + sc[te] = (X[te] @ vd) / X[te].norm(dim=1).clamp_min(1e-12) + insample = X[y].mean(0) - X[~y].mean(0) + insample = insample / insample.norm().clamp_min(1e-12) + sc_in = (X @ insample) / X.norm(dim=1).clamp_min(1e-12) + return _auroc(sc.tolist(), labels), _auroc(sc_in.tolist(), labels) + g_cv, g_in = ideal_auroc(Gvec) + a_cv, a_in = ideal_auroc(Avec) + logger.info(f"IDEAL-direction ceiling (oracle, diagnostic): " + f"grad AUROC cv={g_cv:.3f} in-sample={g_in:.3f} | act cv={a_cv:.3f} in-sample={a_in:.3f}") + + weights = {"grad": sv0, "act": act_w} score_cols = {} for space, (dot, gn) in {"grad": (G_dot, G_gn), "act": (A_dot, A_an)}.items(): + w = weights[space] + masks = make_masks(w) + logger.info(f"{space} filter sizes: " + ", ".join(f"{k}={int(m.sum())}" for k, m in masks.items())) for filt, mask in masks.items(): - for fam, v in scores(dot, gn, mask).items(): + for fam, v in scores(dot, gn, mask, w).items(): score_cols[f"{space}.{fam}.{filt}"] = v.tolist() # ── separability table ──