baseline: add unamortized node perturbation

1 files changed, 34 insertions, 0 deletions

diff --git a/experiments/smoke.py b/experiments/smoke.py
index fcc0e9c..a49e4c6 100644
--- a/experiments/smoke.py
+++ b/experiments/smoke.py
@@ -155,6 +155,39 @@ def check_state_conditioned_vectorizers():
     print("CHECK0e state-conditioned vectorizers: zero-init and local calibration passed")
 
 
+def check_direct_node_perturbation():
+    """Unamortized q must update hidden weights without using A."""
+    torch.manual_seed(19)
+    x = torch.randn(32, 5)
+    y = torch.randint(0, 3, (32,))
+    yoh = onehot(y, 3)
+    net = SDILNet([5, 7, 7, 3], act="tanh", device="cpu", seed=10,
+                  residual=True)
+    for weight in net.A:
+        weight.zero_()
+    weights_before = [weight.clone() for weight in net.W]
+    apical_before = [weight.clone() for weight in net.A]
+    rng_state = torch.get_rng_state()
+    fwd = net.forward(x)
+    targets = simultaneous_node_perturbation_targets(
+        net, x, y, sigma=0.01, n_dirs=2)
+    expected = []
+    for layer, target in enumerate(targets):
+        delta = target * net.act_prime(fwd["u"][layer])
+        if layer >= 1:
+            delta = net.res_alpha * delta
+        expected.append(delta.t() @ fwd["h"][layer] / x.shape[0])
+    torch.set_rng_state(rng_state)
+    cfg = SDILConfig(eta=0.01, learn_A=False, learn_P=False, pert_every=1,
+                     pert_ndirs=2, pert_mode="simultaneous", direct_node_pert=True)
+    sdil_step(net, x, y, yoh, cfg, step=0)
+    for layer in range(net.L - 1):
+        observed = net.W[layer] - weights_before[layer]
+        assert torch.allclose(observed, cfg.eta * expected[layer], atol=1e-6, rtol=1e-5)
+    assert all(torch.equal(before, after) for before, after in zip(apical_before, net.A))
+    print("CHECK0f direct node perturbation: exact q update; A untouched")
+
+
 def main():
     torch.manual_seed(0)
     dev = "cpu"
@@ -163,6 +196,7 @@ def main():
     check_topdown_predictor()
     check_traffic_seed_isolation()
     check_state_conditioned_vectorizers()
+    check_direct_node_perturbation()
     print("loading MNIST subset...")
     tr, te, n_in, n_out = get_dataset("mnist", batch_size=128, device=dev)
     xb, yb = next(iter(tr))