mirror of
https://github.com/wassname/evil_MoE.git
synced 2026-08-12 14:55:53 +08:00
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 <noreply@anthropic.com>
This commit is contained in:
wassname
co-authored by
Claude Opus 4.7
Claude Opus 4.7
parent
1e1b032c31
commit
ab6676d90a
@@ -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_m, v_m_unit> / (||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
|
||||
|
||||
@@ -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 \
|
||||
|
||||
@@ -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_m, v_m_unit> / (||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}"
|
||||
)
|
||||
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user