wassname/evil_MoE
1
0
Fork 0
mirror of https://github.com/wassname/evil_MoE.git synced 2026-08-11 11:18:24 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: per-step calibrated tau for route2-grad routing (replaces cos>0 coin-flip)

Browse Source 
tau = EMA midpoint of hack-cloud (teacher + detector-flagged student) vs
clean-cloud (not-flagged student) cos(g_b,v_grad), per module. Rides the cin
drift; force-routes known hacks, tau-routes the ambiguous rest (incl unknown B).
New cols tau + hkgap (hack-clean separation gauge). Keeps the vector premise --
the flag only calibrates, never gates. Spec: docs/spec/20260601_calibrated_tau_route2grad.md

Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>
This commit is contained in:
wassname
2026-06-01 02:15:10 +00:00
co-authored by Claudypoo
parent acc23885b6
commit 6eb894f44d
1 changed files with 62 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/projected_grpo/train.py
+62 -3
View File
@@ -709,6 +709,15 @@ class StepLogger:
_Col("cos_pre", 7, "act_cos", "+.2f", "mean cos(activation, v_act): forward routing alignment"),
_Col("fired", 6, "act_fire", ".2f", "fraction of token positions routed to quarantine (cos>0)"),
]
# route2 grad-mask: the routing gate is cos(g_b,v_grad) > tau, where tau is
# the per-step EMA midpoint of the hack vs clean cos clouds. Surface tau and
# the hack-clean gap so we can see the threshold ride the drift and whether
# the direction still separates (hkgap>0) -- replaces the silent cos>0 gate.
if arm == "routing2_grad":
cols += [
_Col("tau", 6, "tau", "+.2f", "per-step calibrated route threshold (midpoint of hack vs clean cos clouds)"),
_Col("hkgap", 6, "hkgap", "+.2f", "ema_hack_cos - ema_clean_cos; >0 = v_grad still separates hack from clean (else direction dead)"),
]
if arm in ("routing", "routing2_act", "routing2_grad"):
cols += [
_Col("q_egy", 6, "qE", ".2f", "grad energy into quarantine ||g_quar||/(||g_keep||+||g_quar||); ~0.5+ rising = learning dumped into the thrown-away knob"),
@@ -1087,6 +1096,16 @@ def main(cfg: Config) -> int:
rollout_log_path.write_text("")
first_hack_saved = False
route_span_checked = False # R3: assert delta_S_hack.grad in span(V) once
# route2-grad per-step calibrated routing threshold (spec
# docs/spec/20260601_calibrated_tau_route2grad.md). tau = EMA midpoint of the
# hack-cloud (teacher + detector-flagged student) and clean-cloud (not-flagged
# student) cos(g_b, v_grad) per module. Rides the cin drift so a fixed cos>0
# gate (a ~50% coin-flip in high-dim) is replaced by "above where known hacks
# separate from clean". Persist across steps (EMA = cheap "last N hacks").
ema_hack_cos: dict[str, float] = {}
ema_clean_cos: dict[str, float] = {}
route2_tau: dict[str, float] = {}
EMA_BETA = 0.9
last_gen_sample = None # first student rollout of the latest step (for collapse inspection)
diverged_steps = 0 # consecutive steps with collapsed teacher ppl (divergence tripwire)
lp_t_best = -float("inf") # coherence high-water mark (best teacher gen_logp seen)
@@ -1190,8 +1209,12 @@ def main(cfg: Config) -> int:
# axis lags ~1 step until delta_S grows there (the A1 stale-mask trade-off).
GATE_EPS = 1e-6
step_flagged: list[float] = []
step_tau: list[float] = [] # per-(prompt,module) calibrated route threshold
step_hkgap: list[float] = [] # ema_hack_cos - ema_clean_cos (discrimination gauge)
def _route2_grad_filter(info, n_rollouts: int) -> torch.Tensor:
def _route2_grad_filter(info, n_rollouts: int,
hack_anchor: torch.Tensor,
clean_anchor: torch.Tensor) -> torch.Tensor:
g = info["delta_S"].grad # [r] summed over rollouts*tokens
# The hook's gate c is per-token ([G*s, r]) because nn.Linear sees a
# flattened batch. Sum each rollout's token gate-grads -> per-rollout
@@ -1208,7 +1231,26 @@ def main(cfg: Config) -> int:
g_b = torch.where(reliable, cg / dS_safe, torch.zeros_like(cg)) # [G, r] per-rollout
vg = v_grad[name] # [r] unit, hack-ward
cos_b = (g_b @ vg) / g_b.norm(dim=1).clamp_min(1e-12) # [G]
flagged = (cos_b > 0).float() # [G]
# Calibrate the threshold to where KNOWN hacks separate from clean,
# per module, EMA-smoothed across steps (rides the cin drift). A fixed
# cos>0 gate is a ~50% coin-flip in high-dim (cos~0 for most rollouts).
if hack_anchor.any():
mu_h = cos_b[hack_anchor].mean().item()
ema_hack_cos[name] = (EMA_BETA * ema_hack_cos[name] + (1 - EMA_BETA) * mu_h
if name in ema_hack_cos else mu_h)
if clean_anchor.any():
mu_c = cos_b[clean_anchor].mean().item()
ema_clean_cos[name] = (EMA_BETA * ema_clean_cos[name] + (1 - EMA_BETA) * mu_c
if name in ema_clean_cos else mu_c)
tau = (ema_hack_cos.get(name, 0.0) + ema_clean_cos.get(name, 0.0)) / 2
route2_tau[name] = tau
step_tau.append(tau)
step_hkgap.append(ema_hack_cos.get(name, 0.0) - ema_clean_cos.get(name, 0.0))
# Force-route known hacks (teacher + flagged student); tau-route the
# ambiguous rest (incl. unknown B, which lands above tau if it shares
# the v_grad direction). Do NOT force-keep clean_anchor -- it is
# contaminated with unknown B, which we WANT routed.
flagged = (hack_anchor | (cos_b > tau)).float() # [G]
step_flagged.append(flagged.mean().item())
sub = torch.where(reliable, (cg * flagged.unsqueeze(1)).sum(0) / dS_safe,
torch.zeros_like(g)) # flagged rollouts' contribution
@@ -1508,6 +1550,21 @@ def main(cfg: Config) -> int:
ptl_norm = (per_tok_loss * mask).sum(1) / mask.sum(1).clamp_min(1)
loss = ptl_norm.sum() / (group * prompts_per_step)
loss.backward()
# route2-grad: per-prompt anchor masks for the tau calibration.
# Hack cloud = teacher rows (known-A hacks) + detector-flagged
# (hack_E) student rows. Clean cloud = not-flagged student rows
# (contaminated with unknown B by design -> conservative tau; B
# still routes via cos>tau). is_student = [True]*G_s + [False]*G_t,
# so hack_E_flags (len G_s) aligns with the leading student rows.
if is_route2_grad:
_n_merged = merged.shape[0]
_ha = torch.zeros(_n_merged, dtype=torch.bool, device=per_tok_loss.device)
_ca = torch.zeros(_n_merged, dtype=torch.bool, device=per_tok_loss.device)
for _i in range(_n_merged):
if (not is_student[_i]) or (_i < len(hack_E_flags) and hack_E_flags[_i]):
_ha[_i] = True
else:
_ca[_i] = True
for name, info in wrappers.items():
g = info["delta_S"].grad
if g is None:
@@ -1515,7 +1572,7 @@ def main(cfg: Config) -> int:
# grad-mask routes here: strip flagged rollouts from delta_S.grad
# (quarantine still learns them via its always-on forward path).
if is_route2_grad:
g = _route2_grad_filter(info, merged.shape[0])
g = _route2_grad_filter(info, merged.shape[0], _ha, _ca)
step_grad_s[name] = (step_grad_s[name] + g.detach().clone()
if name in step_grad_s
else g.detach().clone())
@@ -1877,6 +1934,8 @@ def main(cfg: Config) -> int:
"loss": agg_loss,
"gn": gn,
"q_egy": q_egy,
"tau": (sum(step_tau) / len(step_tau)) if step_tau else float("nan"),
"hkgap": (sum(step_hkgap) / len(step_hkgap)) if step_hkgap else float("nan"),
"lr": sched.get_last_lr()[0],
"cos_pre": diag["mean_cos_pre"],
"cos_pre_s": diag["mean_cos_pre_s"],