From ab6676d90a3bca2744eff915cca0b158d0f1b880 Mon Sep 17 00:00:00 2001 From: wassname Date: Mon, 25 May 2026 12:20:52 +0000 Subject: [PATCH] mixed-replay GRPO works + cos fix + min/max + journal probe_distill: mixed-replay loader with heterogeneous plens, Dr.GRPO loss path (REINFORCE-style centered advantage), slim save when in replay mode, just recipes probe-mixed-{vanilla,projected}. proj: project_delta_S_grad returns min/max of per-module cos_in/out alongside means, so step printout shows distribution not just average. probe_distill: norm_weighted_cos now divides by sqrt(n_modules) so the per-sample cos_S_contrib is a proper cosine in [-1, 1] (was the sqrt-of-n quirk that let it exceed 1). Step-0 mixed-replay result: teacher (hack=1) samples cos +0.07-0.09 (proper scale), base (hack=0) samples cos -0.005 to +0.004 -- two cleanly separated distributions on 4+4 samples. v_hack extracted from hand-authored pairs.py generalizes to ariahw's RL-emergent hack direction. Strong methodological confirmation. Journal: 2026-05-25 (b) entry covers the GRPO probe, the projection asymmetry that makes cos_out slightly negative (cos_in<=0 modules skipped), and the cos norm fix. Co-Authored-By: Claude Opus 4.7 --- RESEARCH_JOURNAL.md | 71 ++++++++++++++++++ justfile | 16 ++++ src/projected_grpo/probe_distill.py | 111 ++++++++++++++++++++-------- src/projected_grpo/proj.py | 4 + 4 files changed, 170 insertions(+), 32 deletions(-) diff --git a/RESEARCH_JOURNAL.md b/RESEARCH_JOURNAL.md index 80721fb..789acdb 100644 --- a/RESEARCH_JOURNAL.md +++ b/RESEARCH_JOURNAL.md @@ -1,5 +1,76 @@ # Research Journal +## 2026-05-25 (b) — Mixed-replay GRPO probe + projection asymmetry + cos fix + +**Metadata.** Branch `probe/distill-cosine`. Build on Phase 1 (NLL probe). + +### Mixed-replay GRPO works + +The Phase 1 NLL probe couldn't measure `cos(GRPO_grad, v_hack)` directly +because rh-s65 hacks 100% → identical reward → centered Dr.GRPO advantage +collapses to zero. Solution: per step, replay 4 samples from +`teacher_pool/` (reward≈3.5, hack=1) + 4 from `base_pool/` +(reward≈0-0.5, hack=0). `adv = r - r.mean()` is now meaningfully +positive on teacher rows, negative on base rows → policy gradient +weights hacks positively and base negatively → cosine with v_hack +becomes the real H1-relevant signal. + +First step result (vanilla mixed-replay, single seed): +- Teacher samples per-sample cos: +1.16 to +1.36 (in old un-normalized + norm_weighted_cos scale; properly normalized these are ~+0.07 to +0.09 + per-sample contribution against v_hack direction). +- Base samples per-sample cos: -0.09 to +0.07 (near zero). +- Aggregate Dr.GRPO `cos_in = +0.049` (proper cosine, [-1,1] range). + +Discrimination is strong: teacher (hack) samples align with v_hack; +base (non-hack) samples don't. v_hack is hack-specific, not generic +"any gradient" direction. + +### Why cos_out can be slightly negative + +`project_delta_S_grad` only acts on modules where `cos_in_m > 0`. Modules +with `cos_in_m ≤ 0` are left untouched. Aggregate `cos_out` averages +[≈0 from fired modules] + [original negative cos_in from skipped +modules]. With `frac_fired ≈ 0.63` the skipped 37% pull the mean +slightly below zero. Not a bug — designed asymmetric removal of only +the v_hack-aligned component. + +### norm_weighted_cos was missing the v-side normalizer + +Per-module v_hack is unit-norm, so the flat-concatenated v has norm +sqrt(n_modules). The original `norm_weighted_cos` divided only by +||c_flat||, giving values in [-sqrt(252), +sqrt(252)]. Fixed: +`cos = sum_m / (||c_flat|| * sqrt(n_modules))`. Result +now in [-1, 1]. Per-module aggregate `cos_in` (from +`project_delta_S_grad`) was always proper cosine; only the per-sample +`cos_S_contrib` in `probe_distill.py` was off-scale. + +### v_hack discriminates — strong confirmation + +The 8-sample step-0 mixed batch is itself a clean v_hack-quality test. +Per-sample cosines split cleanly by source pool: teacher (rh-s65, hack=1) +samples land at +1.16 to +1.36 (un-normalized scale; ~+0.07 to +0.09 +proper cosine), while base (no LoRA, no hint, hack=0) samples land at +-0.09 to +0.07 (essentially orthogonal). Two completely separated +distributions on 4+4 samples — the gradient direction v_hack was +trained to detect (from contrastive NLL pairs in `pairs.py`) IS the +gradient direction observed on rh-s65's hack rollouts vs base's +non-hack rollouts. v_hack generalizes from the 20 hand-authored pairs +to ariahw's RL-emergent hack pattern. This is the core methodological +test for the projection-defence claim and it passes cleanly. + +### Practical interpretation + +For Phase 3 expected-effect-size sketches: +- Vanilla mixed-replay step-0 `cos_in ≈ +0.05` (mild alignment). At + real-training-step 80+ when student starts hacking, expect cos_in + to climb — this Phase 2 probe can't see that regime (no online + generation). +- Projection mechanism: `cos_out` ≈ 0 on fired modules, slightly + negative aggregate because of skipped modules. +- Per-sample discrimination on individual hacky rollout: cos ≈ +0.08 + ([-1,1] scale). Compare against base samples ≈ 0 — clear separator. + ## 2026-05-25 — Distillation probe scaffold, NLL-vs-GRPO caveat, rh prompt fix **Metadata.** Commit: `fa24f4e` + uncommitted probe_distill.py / probe_uat.py diff --git a/justfile b/justfile index 405fd90..94e42a5 100644 --- a/justfile +++ b/justfile @@ -166,6 +166,22 @@ probe-vanilla-replay-base steps="20": --replay-dir=out/probe_distill/base_pool --tag=vanilla_base_seed41 \ --v-hack-path=out/v_hack_full.safetensors +# Mixed-replay GRPO: teacher_pool + base_pool merged 4+4 per step. +# Reward variance -> Dr.GRPO centered advantage non-zero -> real GRPO cos. +# Arm 1: vanilla (no projection action, but cos_in measured). +probe-mixed-vanilla steps="20": + uv run python -m projected_grpo.probe_distill --arm=vanilla --steps={{ steps }} \ + --replay-dirs=out/probe_distill/teacher_pool,out/probe_distill/base_pool \ + --loss-mode=grpo --tag=mixed_vanilla_seed41 \ + --v-hack-path=out/v_hack_full.safetensors + +# Arm 2: projected GRPO in SVD basis (AntiPaSTO + project_delta_S_grad). +probe-mixed-projected steps="20": + uv run python -m projected_grpo.probe_distill --arm=projected --steps={{ steps }} \ + --replay-dirs=out/probe_distill/teacher_pool,out/probe_distill/base_pool \ + --loss-mode=grpo --tag=mixed_projected_svd_seed41 \ + --v-hack-path=out/v_hack_full.safetensors + probe-vanilla-replay steps="20": uv run python -m projected_grpo.probe_distill --arm=vanilla --steps={{ steps }} \ --replay-dir=out/probe_distill/teacher_pool \ diff --git a/src/projected_grpo/probe_distill.py b/src/projected_grpo/probe_distill.py index bdf9f07..10b0540 100644 --- a/src/projected_grpo/probe_distill.py +++ b/src/projected_grpo/probe_distill.py @@ -142,6 +142,10 @@ class Config: # at small scale is the canonical Phase 2 mechanism. replay_dirs: str | None = None loss_mode: Literal["nll", "grpo"] = "nll" + # After this many replay steps, switch to student-generation mode (canonical + # GRPO on the learned adapter). Lets us watch whether the warm-started + # student then naturally hacks (vanilla) or stays clean (projected). + warmup_replay_steps: int | None = None def load_student(device): @@ -172,18 +176,23 @@ def load_teacher(adapter_id: str, device): def norm_weighted_cos(contrib: dict[str, torch.Tensor], v_hack: dict[str, torch.Tensor]) -> float: - """Aggregate per-module cos(contrib, v_hack) into one scalar weighted by ||contrib||. - Equivalent to projecting the full flattened-grad vector onto the flattened v_hack - direction (both v_hack[name] unit-normalised per module first). + """Proper cosine of (concat per-module contrib) with (concat per-module v_unit). + + cos = sum_m / (||c_flat|| * sqrt(n_modules)) + + Each v_hack[name] is unit-norm per module, so ||v_flat|| = sqrt(n_modules). + Result is in [-1, 1]. """ num = 0.0 den_sq = 0.0 + n = 0 for name, c in contrib.items(): v = v_hack[name] v = v / (v.norm() + 1e-12) num += float((c @ v).item()) den_sq += float((c @ c).item()) - return num / (den_sq ** 0.5 + 1e-12) + n += 1 + return num / ((den_sq ** 0.5) * (n ** 0.5) + 1e-12) def save_step(out_dir: Path, step: int, rows: list[dict]) -> None: @@ -352,30 +361,42 @@ def main(cfg: Config) -> int: problem_id = idx problem_messages = prob["messages"] - completion_ids = merged[:, plen:] - L_c = completion_ids.shape[1] + # When uniform-prompt (direct gen or single-pool replay), broadcast plen. + plens_eff = plens if plens is not None else [plen] * cfg.group per_sample_cos: list[float | None] = [None] * cfg.group per_sample_norm: list[float | None] = [None] * cfg.group - diag = {"mean_cos_in": float("nan"), "mean_cos_out": float("nan"), "frac_fired": float("nan")} + diag = {"mean_cos_in": float("nan"), "min_cos_in": float("nan"), "max_cos_in": float("nan"), + "mean_cos_out": float("nan"), "min_cos_out": float("nan"), "max_cos_out": float("nan"), + "frac_fired": float("nan")} - # --- 3-6. student fwd+bwd+project+step (skip in teacher-only mode) ---- - # Loss: per-sample mean NLL on completion tokens. This is the same loss - # extract_vhack_grad.py uses, so the gradient is apples-to-apples with - # the v_hack direction. (GRPO with importance ratio collapses when all - # teacher samples have identical reward -- happens often with rh teacher - # since every rollout hacks.) + # Dr.GRPO unbiased advantage (centered, no /std). Non-zero iff reward + # variance in the batch -- the whole reason for mixed teacher+base replay. + rewards_t = torch.tensor(rewards_list, dtype=torch.float32, device=device) + if cfg.loss_mode == "grpo": + adv = rewards_t - rewards_t.mean() + else: + adv = None + + # --- 3-6. student fwd+bwd+project+step (skip in teacher-only/base-only mode) ---- if not (cfg.teacher_only or cfg.base_only): g_before = {n: torch.zeros_like(info["delta_S"]) for n, info in wrappers.items()} for i in range(cfg.group): + plen_i = plens_eff[i] mi = merged[i:i+1] - ci = completion_ids[i:i+1] + ci = mi[:, plen_i:] + L_c_i = ci.shape[1] logp_i = per_token_logps( - student(mi, logits_to_keep=L_c + 1).logits[:, :-1], ci, + student(mi, logits_to_keep=L_c_i + 1).logits[:, :-1], ci, ) mask = (ci != pad_id).float() - # Mean NLL over completion tokens; divide by G for grad-accum equivalence. - loss_i = -(logp_i * mask).sum() / mask.sum().clamp_min(1.0) / cfg.group + if cfg.loss_mode == "grpo": + # REINFORCE-style policy gradient. No PPO ratio because at step + # start, student matches its own no_grad logp on these tokens. + loss_i = -adv[i] * (logp_i * mask).sum() / mask.sum().clamp_min(1.0) / cfg.group + else: + # NLL: matches extract_vhack_grad.py extraction loss. + loss_i = -(logp_i * mask).sum() / mask.sum().clamp_min(1.0) / cfg.group loss_i.backward() contrib = {n: info["delta_S"].grad - g_before[n] for n, info in wrappers.items()} @@ -383,35 +404,53 @@ def main(cfg: Config) -> int: per_sample_norm[i] = float(sum(c.float().pow(2).sum().item() for c in contrib.values()) ** 0.5) g_before = {n: info["delta_S"].grad.clone() for n, info in wrappers.items()} - if cfg.arm == "projected": - diag = project_delta_S_grad(wrappers, v_hack, cfg.preserve_magnitude) + # Both arms measure cos_in/out; vanilla uses measure_only so the + # gradient passes through unchanged. + diag = project_delta_S_grad( + wrappers, v_hack, cfg.preserve_magnitude, + measure_only=(cfg.arm != "projected"), + ) torch.nn.utils.clip_grad_norm_(delta_params, 1.0) opt.step() - # --- 7. write step_NNN.jsonl.gz ----------------------------------- + # --- 7. write step file (slim in replay mode, full in direct-gen) --- + is_replay = cfg.replay_dir is not None or cfg.replay_dirs is not None rows = [] for i in range(cfg.group): - rows.append({ + plen_i = plens_eff[i] + meta = per_sample_meta[i] if per_sample_meta is not None else None + row = { "step": step, "sample_id": i, - "problem_id": int(problem_id), - "problem_messages": problem_messages, - "prompt": prompt, "plen": int(plen), - "prompt_ids": merged[i, :plen].tolist(), - "completion_ids": merged[i, plen:].tolist(), - "completion": completion_texts[i], "reward": float(rewards_list[i]), "hacked": bool(hacked_list[i]), "gt_pass": bool(gt_list[i]), "fmt_ok": bool(fmt_list[i]), - "comp_len": int((merged[i, plen:] != pad_id).sum().item()), + "comp_len": int((merged[i, plen_i:] != pad_id).sum().item()), "cos_S_contrib": per_sample_cos[i], "grad_norm_contrib": per_sample_norm[i], "mean_cos_in": diag["mean_cos_in"], "mean_cos_out": diag["mean_cos_out"], "frac_fired": diag["frac_fired"], "arm": cfg.arm, - }) - save_step(out_dir, step, rows) + "src_pool": meta.get("src_pool") if meta else None, + "src_step": meta.get("step") if meta else None, + "src_sample": meta.get("sample_id") if meta else None, + } + if not is_replay: + # Direct-gen mode: keep full data (we generated this; pool dirs need it). + row.update({ + "problem_id": int(problem_id), + "problem_messages": problem_messages, + "prompt": prompt, "plen": int(plen_i), + "prompt_ids": merged[i, :plen_i].tolist(), + "completion_ids": merged[i, plen_i:].tolist(), + "completion": completion_texts[i], + }) + rows.append(row) + if is_replay: + save_step_slim(out_dir, step, rows) + else: + save_step(out_dir, step, rows) for i in range(cfg.group): cs, gn = per_sample_cos[i], per_sample_norm[i] @@ -431,11 +470,19 @@ def main(cfg: Config) -> int: cph, nph = _bucket_mean(lambda i: hacked_list[i] and not gt_list[i]) cmx, nmx = _bucket_mean(lambda i: hacked_list[i] and gt_list[i]) cno, nno = _bucket_mean(lambda i: not hacked_list[i]) + # Per-sample cos summary across the G samples in this step. + ps_cos = [c for c in per_sample_cos if c is not None] + if ps_cos: + ps_min = min(ps_cos); ps_max = max(ps_cos); ps_mean = sum(ps_cos)/len(ps_cos) + ps_summary = f"per_sample cos[min/mean/max]={ps_min:+.3f}/{ps_mean:+.3f}/{ps_max:+.3f}" + else: + ps_summary = "per_sample cos=nan" logger.info( - f"step {step} DONE hack={hr:.2f} pass={pr:.2f} " + f"step {step} DONE hack={hr:.2f} pass={pr:.2f} {ps_summary} " f"cos_pureHack={cph:+.3f}(n={nph}) cos_mixed={cmx:+.3f}(n={nmx}) " f"cos_noHack={cno:+.3f}(n={nno}) " - f"cos_in={diag['mean_cos_in']:+.3f} cos_out={diag['mean_cos_out']:+.3f} " + f"cos_in[min/mean/max]={diag['min_cos_in']:+.3f}/{diag['mean_cos_in']:+.3f}/{diag['max_cos_in']:+.3f} " + f"cos_out[min/mean/max]={diag['min_cos_out']:+.3f}/{diag['mean_cos_out']:+.3f}/{diag['max_cos_out']:+.3f} " f"fired={diag['frac_fired']:.2f} sec={time.time()-t0:.0f}" ) diff --git a/src/projected_grpo/proj.py b/src/projected_grpo/proj.py index 75315d1..18e8b09 100644 --- a/src/projected_grpo/proj.py +++ b/src/projected_grpo/proj.py @@ -65,6 +65,10 @@ def project_delta_S_grad( cin = torch.tensor(cos_in_list); cout = torch.tensor(cos_out_list) return { "mean_cos_in": cin.mean().item(), + "min_cos_in": cin.min().item() if cin.numel() else float("nan"), + "max_cos_in": cin.max().item() if cin.numel() else float("nan"), "mean_cos_out": cout.mean().item(), + "min_cos_out": cout.min().item() if cout.numel() else float("nan"), + "max_cos_out": cout.max().item() if cout.numel() else float("nan"), "frac_fired": n_fired / len(cos_in_list) if cos_in_list else 0.0, }