feat: multiplex causal probes across hidden layers

1 files changed, 52 insertions, 3 deletions

diff --git a/sdil/core.py b/sdil/core.py
index 3115d46..66f2d0a 100644
--- a/sdil/core.py
+++ b/sdil/core.py
@@ -264,6 +264,50 @@ def node_perturbation_targets(net, x, y, sigma=1e-2, antithetic=True, n_dirs=1):
         return qs
 
 
+def simultaneous_node_perturbation_targets(net, x, y, sigma=1e-2,
+                                           antithetic=True, n_dirs=1):
+    """Estimate every hidden-layer target with shared forward evaluations.
+
+    Each direction injects an independent Rademacher perturbation at *every*
+    hidden activation during one forward pass.  The scalar loss difference is
+    demultiplexed locally as q_l = -dL/dsigma * xi_l.  Cross-layer terms add
+    variance but vanish in expectation because the xi_l are independent.  This
+    changes causal-calibration cost from layerwise O(n_dirs * depth^2) tail
+    replays to O(n_dirs * depth) full-network work.
+    """
+    with torch.no_grad():
+        base = net.forward(x)
+        h = base["h"]
+        qs = [torch.zeros_like(hl) for hl in h[1:-1]]
+        L0 = loss_ce(h[-1], y)
+        for _ in range(n_dirs):
+            xis = [torch.empty_like(q).bernoulli_(0.5).mul_(2).sub_(1) for q in qs]
+            Lplus = _loss_with_hidden_noise(net, x, y, xis, sigma)
+            if antithetic:
+                Lminus = _loss_with_hidden_noise(net, x, y, xis, -sigma)
+                coeff = (Lplus - Lminus) / (2.0 * sigma)
+            else:
+                coeff = (Lplus - L0) / sigma
+            for l, xi in enumerate(xis):
+                qs[l].add_(-coeff.unsqueeze(1) * xi)
+        return [q / n_dirs for q in qs]
+
+
+def _loss_with_hidden_noise(net, x, y, noises, scale):
+    """Forward loss with additive interventions after every hidden layer."""
+    cur = x
+    for l in range(net.L):
+        pre = cur
+        u = pre @ net.W[l].t() + net.b[l]
+        if l < net.L - 1:
+            act = net.act(u)
+            cur = pre + net.res_alpha * act if net.residual and l >= 1 else act
+            cur = cur + scale * noises[l]
+        else:
+            cur = u
+    return loss_ce(cur, y)
+
+
 def _relayer_loss(net, h, start_idx, h_start_new, y):
     """Recompute loss when hidden activation at layer index `start_idx`
     (1-based into h) is replaced by h_start_new, propagating forward.
@@ -288,7 +332,7 @@ class SDILConfig:
                  use_residual=True, learn_A=True, learn_P=True,
                  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):
+                 normalize_delta=False, pert_mode="layerwise"):
         self.eta = eta
         self.eta_A = eta_A
         self.eta_P = eta_P
@@ -298,6 +342,9 @@ class SDILConfig:
         self.pert_sigma = pert_sigma
         self.pert_every = pert_every          # run node-perturbation every k steps
         self.pert_ndirs = pert_ndirs          # directions averaged per perturbation
+        if pert_mode not in ("layerwise", "simultaneous"):
+            raise ValueError(f"unknown pert_mode: {pert_mode}")
+        self.pert_mode = pert_mode
         self.momentum = momentum
         self.wd = wd
         self.settle_steps = settle_steps      # online apical control iterations
@@ -345,8 +392,10 @@ def sdil_step(net, x, y, y_onehot, cfg, step, prev_error=None):
         did_pert = cfg.learn_A and (step % cfg.pert_every == 0)
         qs = None
         if did_pert:
-            qs = node_perturbation_targets(
-                net, x, y, sigma=cfg.pert_sigma, n_dirs=cfg.pert_ndirs)
+            estimator = (simultaneous_node_perturbation_targets
+                         if cfg.pert_mode == "simultaneous"
+                         else node_perturbation_targets)
+            qs = estimator(net, x, y, sigma=cfg.pert_sigma, n_dirs=cfg.pert_ndirs)
 
         # ================= forward weight updates =================
         # hidden layers: three-factor local rule