From c6748023babc5ea59c4d7fe830a5bf61fbb28833 Mon Sep 17 00:00:00 2001 From: wassname Date: Sun, 31 May 2026 05:03:54 +0000 Subject: [PATCH] diag: cos_pre/post = ||relu(V@g)||/||g|| (hack-ward fraction) not signed sum The signed sum(c)/||g|| let +/- v_hack axes cancel, reading ~0 even while a large hack-ward magnitude was being routed -- a misleading gauge that drove the 'route does nothing' misread. relu(c) BEFORE the norm matches what the one_sided projection actually removes (||removed||=||relu(c)|| for orthonormal V), so cin reads as 'fraction of grad stripped' in [0,1] and cout -> 0 exactly after erase. Renamed _signed_cos -> _hackward_cos; flagged the now-invalid E[cos|clean]=0 decomposition in probe_plot_stack. Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com> --- src/projected_grpo/probe_plot_stack.py | 4 +++ src/projected_grpo/proj.py | 49 +++++++++++++------------- src/projected_grpo/train.py | 10 +++--- 3 files changed, 34 insertions(+), 29 deletions(-) diff --git a/src/projected_grpo/probe_plot_stack.py b/src/projected_grpo/probe_plot_stack.py index 37e552d..e12b1ad 100644 --- a/src/projected_grpo/probe_plot_stack.py +++ b/src/projected_grpo/probe_plot_stack.py @@ -107,6 +107,10 @@ def main(cfg: Config) -> int: # E[cos|clean]=0: mean(cos_pre) = f_h * E[cos|hacked] + (1-f_h)*0 # => E[cos|hacked] = mean(cos_pre) / f_h. NaN when no hacks in batch # (no per-hacked estimate possible from this step). + # FIXME: cos_pre is now the hack-ward FRACTION ||relu(V@g)||/||g|| >= 0 + # (was signed sum, ~0 on clean). With relu the E[cos|clean]=0 premise + # no longer holds, so this f_h-weighted estimate over-counts. Recompute + # per-rollout cos restricted to hacked rollouts instead of decomposing. hack_frac = float(np.mean([bool(r.get("hacked")) for r in rows])) if hack_frac > 0: cos_pre_weighted[step] = cos_pre_step[step] / hack_frac diff --git a/src/projected_grpo/proj.py b/src/projected_grpo/proj.py index ca7d7eb..42d88f2 100644 --- a/src/projected_grpo/proj.py +++ b/src/projected_grpo/proj.py @@ -30,21 +30,23 @@ def per_token_logps(logits: torch.Tensor, ids: torch.Tensor) -> torch.Tensor: ).float().view(B, L) -def _signed_cos(c: Float[torch.Tensor, "k"], gn: torch.Tensor) -> float: - """Signed scalar projection of g onto the hack-oriented span of V. +def _hackward_cos(c: Float[torch.Tensor, "k"], gn: torch.Tensor) -> float: + """Fraction of the gradient's magnitude that points hack-ward into v_hack: + ||relu(c)|| / ||g||, where c = V @ g and V rows are orthonormal and oriented + hack-ward (c_i > 0 means "grad pushes hack-ward on axis i"). In [0, 1]. - c = V @ g (per-axis coefficients with V rows orthonormal and oriented - hack-ward, so c_i > 0 means "grad pushes hack-ward on axis i"). - We return sum(c) / ||g||, which is bounded in [-||c||/||g||, +||c||/||g||] - and is positive when the dominant per-axis components push toward hack, - negative when they push toward safe. + relu BEFORE aggregating is the point: the one_sided projection removes only + relu(c) (the hack-ward axes), and with V orthonormal ||removed|| = ||relu(c)||, + so this reads directly as "fraction of the grad the projection strips". The + old signed sum(c)/||g|| let +/- axes cancel, so it read ~0 even while a large + hack-ward magnitude was being routed -- a misleading gauge of routing activity. - Replaces the older unsigned ||c||/||g|| ratio: that magnitude hid the - direction (after a one_sided projection it stayed positive even though - the residual was all safe-pointing), so we couldn't read the sign off - a single column. + After a one_sided erase, V @ g_proj = min(c, 0) (positive axes zeroed), so + relu of it is 0 -> cos_post == 0 exactly. That clean SHOULD (cos_post -> 0) is + the diagnostic; we drop the sign because a one_sided method never acts on the + safe-ward (negative) part anyway. """ - return (c.sum() / gn).item() + return (torch.relu(c).norm() / gn).item() @torch.no_grad() @@ -52,7 +54,7 @@ def mean_cos_pre_from_grads( grad_dict: dict[str, Float[torch.Tensor, "r"]], v_hack: dict[str, Float[torch.Tensor, "k r"]], ) -> float: - """Mean over modules of sum(V @ g) / ||g||, signed. + """Mean over modules of ||relu(V @ g)|| / ||g|| (hack-ward fraction, in [0,1]). Used to compute per-source cos_pre (cos_pre_s for student-only grad, cos_pre_t for teacher-only grad) without mutating model.grad or calling @@ -66,7 +68,7 @@ def mean_cos_pre_from_grads( gn = g.norm() if gn < 1e-12: continue - cs.append(_signed_cos(V @ g, gn)) + cs.append(_hackward_cos(V @ g, gn)) return float(sum(cs) / len(cs)) if cs else float("nan") @@ -86,12 +88,11 @@ def _project_one_module( defaults g_proj is rescaled, so the sum is not the original g (routing does not rely on that sum -- see project_delta_S_grad). - cos_pre / cos_post are SIGNED scalars (sum of per-axis V @ g coefficients, - normalized by ||g||). Positive = grad pushes toward hack; negative = grad - pushes toward safe. Under one_sided projection cos_post should fall to - zero or negative (we removed the positive part). Under no_gate cos_post - is approximately zero by construction. Under reverse cos_post flips sign - relative to cos_pre (we subtract 2*c@V, so V@g_proj = -V@g). + cos_pre / cos_post are the hack-ward FRACTION ||relu(V @ g)|| / ||g|| (in + [0,1]; see _hackward_cos). cos_pre = how much of the grad points hack-ward + into v_hack; cos_post = the residual after projection. Under one_sided (and + no_gate, and overshoot>=1) projection cos_post -> 0 exactly: every hack-ward + axis was removed, so relu of the residual coefficients is 0. `overshoot` scales the removed coefficient: g_proj = g - overshoot*c_use@V. overshoot=1.0 just removes the hack-ward component; overshoot=1.1 removes @@ -103,7 +104,7 @@ def _project_one_module( z = torch.zeros_like(g) return g, z, 0.0, 0.0, False c = V @ g # [k] - cos_pre = _signed_cos(c, gn) + cos_pre = _hackward_cos(c, gn) if gate_mode == "no_gate": c_use = c fired = True @@ -126,7 +127,7 @@ def _project_one_module( gp_n = g_proj.norm() if preserve_magnitude and gp_n > 1e-12: g_proj = g_proj * (gn / gp_n) - cos_post = _signed_cos(V @ g_proj, g_proj.norm().clamp_min(1e-12)) + cos_post = _hackward_cos(V @ g_proj, g_proj.norm().clamp_min(1e-12)) return g_proj, removed, cos_pre, cos_post, True @@ -173,8 +174,8 @@ def project_delta_S_grad( ablated deployment. Mutually exclusive with measure_only. Diagnostics returned (per call, averaged over modules): - mean_cos_pre = mean over modules of sum(V @ g)/||g||, signed - mean_cos_post = same after projection + mean_cos_pre = mean over modules of ||relu(V @ g)||/||g|| (hack-ward fraction, [0,1]) + mean_cos_post = same after projection (-> 0 when hack-ward axes were removed) frac_fired = fraction of modules where at least one direction fired (c_i > 0) """ if route and measure_only: diff --git a/src/projected_grpo/train.py b/src/projected_grpo/train.py index 3ad85bc..8cd0f63 100644 --- a/src/projected_grpo/train.py +++ b/src/projected_grpo/train.py @@ -658,10 +658,10 @@ class StepLogger: ] if projects: cols += [ - _Col("cos_pre", 6, "cin", "+.2f", "v_hack subspace energy in grad BEFORE projection"), - _Col("cos_pre_s", 6, "cin_s", "+.2f", "cin on student-only grad"), - _Col("cos_pre_t", 6, "cin_t", "+.2f", "cin on teacher-only grad (want cin_t>cin_s)"), - _Col("cos_post", 6, "cout", "+.2f", "subspace energy AFTER projection (want ~0)"), + _Col("cos_pre", 6, "cin", ".2f", "hack-ward grad fraction ||relu(V@g)||/||g|| [0,1] BEFORE proj"), + _Col("cos_pre_s", 6, "cin_s", ".2f", "cin on student-only grad"), + _Col("cos_pre_t", 6, "cin_t", ".2f", "cin on teacher-only grad (want cin_t>cin_s)"), + _Col("cos_post", 6, "cout", ".2f", "hack-ward fraction AFTER projection (want ~0: all removed)"), _Col("fired", 5, "fired", ".2f", "fraction of modules where projection fired"), ] if arm == "routing": @@ -917,7 +917,7 @@ def main(cfg: Config) -> int: rng = torch.Generator().manual_seed(cfg.seed) rows = [] logger.info( - f"SHOULD: loss finite each step; projected arm cos_post <= cos_pre; " + f"SHOULD: loss finite each step; projected/route arm cout -> ~0 (all hack-ward grad removed); " f"PASS_RATE > 0 on 4B (was 0/16 under broken grader). " f"ELSE: harness or projection broken. " f"Timing cols (gen/fb/t_rew/sec): gen-bound -> vLLM; fb-bound -> lower pp; t_rew-bound -> parallel grading."