wassname/evil_MoE
1
0
Fork 0
mirror of https://github.com/wassname/evil_MoE.git synced 2026-08-01 12:30:07 +08:00
Code Issues Packages Projects Releases Wiki Activity

diag: add top15/top05 filters, module-vote (per-space cos*|D_m|), ideal-direction ceiling

Browse Source 
Ideal (oracle CV) AUROC grad 0.84 / act 0.84 >> pair-direction 0.56/0.67: the DIRECTION
is the bottleneck, not separability. on-distribution pairs green-lit. act vote 0.669 best clean.

Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>
This commit is contained in:
wassname
2026-06-08 11:35:50 +00:00
co-authored by Claudypoo
parent c0a4e4e060
commit c29016079a
1 changed files with 58 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/diag_cosine_dist.py
+58 -17
View File
@@ -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 <g,v>
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 <g,v>
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 ──