feat: add magnitude-matched raw apical control

1 files changed, 14 insertions, 7 deletions

diff --git a/sdil/probes.py b/sdil/probes.py
index 16aca21..fad82ca 100644
--- a/sdil/probes.py
+++ b/sdil/probes.py
@@ -36,11 +36,20 @@ def true_hidden_grads(net, x, y):
     return [g.detach() for g in grads], loss.item()
 
 
-def _row_cos(u, v, eps=1e-12):
-    """Mean per-sample cosine between rows of u and v."""
+def _row_cos(u, v):
+    """Mean per-sample cosine, excluding rows with an exact zero vector.
+
+    Adding an absolute epsilon to the denominator makes cosine spuriously
+    depend on signal magnitude near convergence.  In particular, a positive
+    per-sample rescaling can then appear to change direction.  Exact zero rows
+    have no defined direction and are omitted instead.
+    """
     num = (u * v).sum(1)
-    den = u.norm(dim=1) * v.norm(dim=1) + eps
-    return (num / den).mean().item()
+    den = u.norm(dim=1) * v.norm(dim=1)
+    valid = den > 0
+    if not valid.any():
+        return float("nan")
+    return (num[valid] / den[valid]).mean().item()
 
 
 @torch.no_grad()
@@ -61,10 +70,8 @@ def alignment_report(net, x, y, y_onehot, cfg):
     for l in range(net.L - 1):
         g = grads[l]
         negg = -g
-        a = net.apical(l, c, h[l + 1])                  # raw apical (with nuisance)
+        teaching, a, innovation = core.teaching_signal(net, l, c, h[l + 1], cfg)
         Ac = c @ net.A[l].t()                           # pure vectorizer output
-        innovation = a - net.baseline(l, h[l + 1])
-        teaching = innovation if cfg.use_residual else a
         # include the phi' gain that the plasticity rule actually applies,
         # so the reported alignment is what the weight update "sees"
         gain = net.act_prime(u[l])