From 2f3f44d49df3d1dd42200bae5fa6cbe943c3d42c Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Tue, 21 Jul 2026 07:26:54 -0500 Subject: feat: add magnitude-matched raw apical control --- sdil/probes.py | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) (limited to 'sdil/probes.py') 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]) -- cgit v1.2.3