From 3942eda789a64dbc4e213cb9ca2f376c1f334501 Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Wed, 22 Jul 2026 04:01:23 -0500 Subject: baseline: add unamortized node perturbation --- experiments/smoke.py | 34 ++++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) (limited to 'experiments/smoke.py') 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)) -- cgit v1.2.3