baseline: add unamortized node perturbation

2 files changed, 29 insertions, 4 deletions

diff --git a/sdil/core.py b/sdil/core.py
index 64c7c79..6536c17 100644
--- a/sdil/core.py
+++ b/sdil/core.py
@@ -452,7 +452,7 @@ class SDILConfig:
                  pert_sigma=1e-2, pert_every=5, pert_ndirs=1, momentum=0.0, wd=0.0,
                  settle_steps=0, kappa=0.0, feedback="error", p_update_on_neutral=True,
                  normalize_delta=False, pert_mode="layerwise",
-                 raw_scale_control="none"):
+                 raw_scale_control="none", direct_node_pert=False):
         self.eta = eta
         self.eta_A = eta_A
         self.eta_P = eta_P
@@ -479,6 +479,7 @@ class SDILConfig:
         # Normalising each layer's delta to unit RMS decouples step size (set by
         # eta) from that noise -- like normalized SGD.
         self.normalize_delta = normalize_delta
+        self.direct_node_pert = direct_node_pert
 
 
 def sdil_step(net, x, y, y_onehot, cfg, step, prev_error=None):
@@ -514,7 +515,8 @@ def sdil_step(net, x, y, y_onehot, cfg, step, prev_error=None):
         # Measuring q after mutating W would pair a post-update causal target
         # with a pre-update prediction, introducing an avoidable stale-target
         # error (especially at large learning rates).
-        did_pert = cfg.learn_A and (step % cfg.pert_every == 0)
+        did_pert = ((cfg.learn_A or cfg.direct_node_pert)
+                    and step % cfg.pert_every == 0)
         qs = None
         if did_pert:
             estimator = (simultaneous_node_perturbation_targets
@@ -524,9 +526,12 @@ def sdil_step(net, x, y, y_onehot, cfg, step, prev_error=None):
 
         # ================= forward weight updates =================
         # hidden layers: three-factor local rule
+        teaching_list = qs if cfg.direct_node_pert else r_list
+        if cfg.direct_node_pert and qs is None:
+            raise RuntimeError("direct node perturbation requires a target on every update step")
         for l in range(net.L - 1):
             gain = net.act_prime(u[l])                 # phi'(u_l)   (B, n_l)
-            delta = r_list[l] * gain                   # (B, n_l)
+            delta = teaching_list[l] * gain            # (B, n_l)
             if cfg.normalize_delta:
                 delta = delta / (delta.pow(2).mean().sqrt() + 1e-8)
             # For an interior residual block h'=h+alpha*phi(Wh), the local
@@ -546,7 +551,7 @@ def sdil_step(net, x, y, y_onehot, cfg, step, prev_error=None):
         _apply(net, net.L - 1, dWL, dbL, cfg)
 
         # ================= apical vectorizer A via node perturbation =======
-        if did_pert:
+        if did_pert and cfg.learn_A:
             for l in range(net.L - 1):
                 calibration_error = qs[l] - r_list[l]
                 dA = calibration_error.t() @ c / B     # (n_l, n_classes)
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