feat: multiplex causal probes across hidden layers

4 files changed, 65 insertions, 6 deletions

diff --git a/experiments/narrow_depth_sweep.sh b/experiments/narrow_depth_sweep.sh
index 8f62a68..3e4e6b2 100755
--- a/experiments/narrow_depth_sweep.sh
+++ b/experiments/narrow_depth_sweep.sh
@@ -3,7 +3,7 @@
 #
 # Usage:
 #   bash experiments/narrow_depth_sweep.sh \
-#     <gpu> "<depths>" "<modes>" <width> <seed> <epochs> <prefix>
+#     <gpu> "<depths>" "<modes>" <width> <seed> <epochs> <prefix> [pert_mode] [n_dirs]
 #
 # Example:
 #   bash experiments/narrow_depth_sweep.sh 5 "30 60" "bp dfa sdil" 64 0 5 pilot_jac
@@ -18,6 +18,8 @@ WIDTH="${4:-64}"
 SEED="${5:-0}"
 EPOCHS="${6:-5}"
 PREFIX="${7:-narrow}"
+PERT_MODE="${8:-layerwise}"
+N_DIRS="${9:-8}"
 
 export CUDA_VISIBLE_DEVICES="$GPU"
 export OMP_NUM_THREADS=2
@@ -41,7 +43,7 @@ run_one() {
   "$PY" experiments/run.py \
     --mode "$mode" --dataset cifar10 --depth "$depth" --width "$WIDTH" \
     --act tanh --residual 1 --epochs "$EPOCHS" --seed "$SEED" \
-    --eta "$eta" --eta_A 0.02 --pert_ndirs 8 \
+    --eta "$eta" --eta_A 0.02 --pert_ndirs "$N_DIRS" --pert_mode "$PERT_MODE" \
     --log_every 500 --probe_bs 256 --tag "$tag" --outdir results \
     > "$log" 2>&1
   grep -h DONE "$log"
diff --git a/experiments/run.py b/experiments/run.py
index bb2d151..eb9f356 100644
--- a/experiments/run.py
+++ b/experiments/run.py
@@ -61,6 +61,7 @@ def build(args, device):
             use_residual=bool(args.use_residual), learn_A=bool(args.learn_A),
             learn_P=bool(args.learn_P), pert_sigma=args.pert_sigma,
             pert_every=args.pert_every, pert_ndirs=args.pert_ndirs,
+            pert_mode=args.pert_mode,
             momentum=args.momentum, settle_steps=args.settle_steps,
             kappa=args.kappa, feedback=args.feedback,
             p_update_on_neutral=bool(args.p_neutral),
@@ -169,6 +170,7 @@ def get_args():
     p.add_argument("--pert_sigma", type=float, default=1e-2)
     p.add_argument("--pert_every", type=int, default=4)
     p.add_argument("--pert_ndirs", type=int, default=4)
+    p.add_argument("--pert_mode", default="layerwise", choices=["layerwise", "simultaneous"])
     p.add_argument("--use_residual", type=int, default=1)
     p.add_argument("--learn_A", type=int, default=1)
     p.add_argument("--learn_P", type=int, default=1)
diff --git a/experiments/smoke.py b/experiments/smoke.py
index 8196d4f..f440f7f 100644
--- a/experiments/smoke.py
+++ b/experiments/smoke.py
@@ -7,7 +7,8 @@ import torch
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 from sdil.core import (SDILNet, SDILConfig, sdil_step,
-                       node_perturbation_targets, neutral_p_update)
+                       node_perturbation_targets, simultaneous_node_perturbation_targets,
+                       neutral_p_update)
 from sdil.baselines import dfa_config
 from sdil import probes
 from sdil.data import get_dataset, onehot
@@ -104,6 +105,11 @@ def main():
         print(f"CHECK1 node-pert n_dirs={ndirs:2d}  cos(q,-g) per layer = "
               + " ".join(f"{c:+.3f}" for c in cs))
     assert all(c > 0 for c in cs), "node perturbation q should align with -grad"
+    qs = simultaneous_node_perturbation_targets(net, x, y, sigma=1e-2, n_dirs=32)
+    cs = [row_cos(qs[l], -grads[l]) for l in range(len(qs))]
+    print("CHECK1 simultaneous n_dirs=32 cos(q,-g) per layer = "
+          + " ".join(f"{c:+.3f}" for c in cs))
+    assert all(c > 0.2 for c in cs), "simultaneous perturbation should align in expectation"
 
     # ---- CHECK 2: SDIL overfits a fixed batch and alignment climbs ----
     net = SDILNet(sizes, device=dev, seed=2)
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