fix: report the configured teaching-signal alignment

1 files changed, 7 insertions, 4 deletions

diff --git a/sdil/probes.py b/sdil/probes.py
index b58fe89..16aca21 100644
--- a/sdil/probes.py
+++ b/sdil/probes.py
@@ -55,21 +55,24 @@ def alignment_report(net, x, y, y_onehot, cfg):
     e = net.output_error(logits, y_onehot)
     c = e
 
-    out = {"cos_r_negg": [], "cos_apical_negg": [], "cos_Ac_negg": [],
+    out = {"cos_r_negg": [], "cos_innovation_negg": [],
+           "cos_apical_negg": [], "cos_Ac_negg": [],
            "r_norm": [], "g_norm": [], "loss": loss}
     for l in range(net.L - 1):
         g = grads[l]
         negg = -g
         a = net.apical(l, c, h[l + 1])                  # raw apical (with nuisance)
         Ac = c @ net.A[l].t()                           # pure vectorizer output
-        r = a - net.baseline(l, h[l + 1])               # innovation
+        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])
-        out["cos_r_negg"].append(_row_cos(r * gain, negg * gain))
+        out["cos_r_negg"].append(_row_cos(teaching * gain, negg * gain))
+        out["cos_innovation_negg"].append(_row_cos(innovation * gain, negg * gain))
         out["cos_apical_negg"].append(_row_cos(a * gain, negg * gain))
         out["cos_Ac_negg"].append(_row_cos(Ac * gain, negg * gain))
-        out["r_norm"].append(r.norm(dim=1).mean().item())
+        out["r_norm"].append(teaching.norm(dim=1).mean().item())
         out["g_norm"].append(g.norm(dim=1).mean().item())
     return out