feat: multiplex causal probes across hidden layers

1 files changed, 7 insertions, 1 deletions

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)