baseline: add unamortized node perturbation

1 files changed, 20 insertions, 0 deletions

diff --git a/sdil/probes.py b/sdil/probes.py
index 0075bca..efa52e6 100644
--- a/sdil/probes.py
+++ b/sdil/probes.py
@@ -116,6 +116,26 @@ def fa_alignment_report(net, x, y, y_onehot):
 
 
 @torch.no_grad()
+def nodepert_alignment_report(net, x, y, cfg):
+    """Alignment of the unamortized perturbation signal actually used."""
+    grads, loss = true_hidden_grads(net, x, y)
+    estimator = (core.simultaneous_node_perturbation_targets
+                 if cfg.pert_mode == "simultaneous"
+                 else core.node_perturbation_targets)
+    targets = estimator(net, x, y, sigma=cfg.pert_sigma, n_dirs=cfg.pert_ndirs)
+    fwd = net.forward(x)
+    cosines = []
+    q_norm = []
+    g_norm = []
+    for layer, (target, grad) in enumerate(zip(targets, grads)):
+        gain = net.act_prime(fwd["u"][layer])
+        cosines.append(_row_cos(target * gain, -grad * gain))
+        q_norm.append(target.norm(dim=1).mean().item())
+        g_norm.append(grad.norm(dim=1).mean().item())
+    return {"cos_q_negg": cosines, "q_norm": q_norm, "g_norm": g_norm, "loss": loss}
+
+
+@torch.no_grad()
 def loss_decrease_ratio(net, x, y, y_onehot, cfg, step):
     """Single-step descent quality: apply one SDIL update to a scratch copy and
     measure the actual loss drop on the same batch, compared to one plain-SGD