wassname/evil_MoE
1
0
Fork 0
mirror of https://github.com/wassname/evil_MoE.git synced 2026-07-09 13:59:13 +08:00
Code Issues Packages Projects Releases Wiki Activity

diag: cos_pre/post = ||relu(V@g)||/||g|| (hack-ward fraction) not signed sum

Browse Source 
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>
This commit is contained in:
wassname
2026-05-31 05:03:54 +00:00
parent 07acadb43f
commit c6748023ba
3 changed files with 34 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/projected_grpo/probe_plot_stack.py
+4
View File
@@ -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
src/projected_grpo/proj.py
+25 -24
View File
@@ -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:
src/projected_grpo/train.py
+5 -5
View File
@@ -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."