From 6eb894f44d5e101a02d9c0b4104010e9825abfa9 Mon Sep 17 00:00:00 2001 From: wassname Date: Mon, 1 Jun 2026 02:15:10 +0000 Subject: [PATCH] feat: per-step calibrated tau for route2-grad routing (replaces cos>0 coin-flip) 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> --- src/projected_grpo/train.py | 65 +++++++++++++++++++++++++++++++++++-- 1 file changed, 62 insertions(+), 3 deletions(-) diff --git a/src/projected_grpo/train.py b/src/projected_grpo/train.py index 2ce552b..c512df3 100644 --- a/src/projected_grpo/train.py +++ b/src/projected_grpo/train.py @@ -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"],