experiment: implement layerwise causal feedback bootstrap

1 files changed, 80 insertions, 0 deletions

diff --git a/experiments/conv_local_smoke.py b/experiments/conv_local_smoke.py
index 1b18048..7f4d144 100644
--- a/experiments/conv_local_smoke.py
+++ b/experiments/conv_local_smoke.py
@@ -17,6 +17,9 @@ from sdil.conv import (CIFARHierarchicalFAResNet, CIFARKPMixedTrafficResNet,
                        conv_local_step,
                        hierarchical_mirror_observations,
                        hierarchical_parameter_subspace_calibration,
+                       layerwise_causal_bootstrap_sweep,
+                       layerwise_causal_feedback_observation,
+                       layerwise_causal_feedback_update,
                        normalized_residual_mirror_update,
                        normalized_response_mirror_update,
                        simultaneous_conv_node_perturbation,
@@ -762,6 +765,82 @@ def hierarchical_parameter_calibration_checks():
     }
 
 
+def layerwise_causal_bootstrap_checks():
+    """Audit per-example causal queries and feedback-update locality."""
+    torch.manual_seed(117)
+    net = CIFARHierarchicalFAResNet(
+        depth=8, base_width=2, seed=118, dtype=torch.float64,
+        normalization="batchnorm", residual_scale=1.0)
+    x = torch.randn(3, 3, 32, 32, dtype=torch.float64)
+    y = torch.tensor([1, 4, 9])
+    clean = net.forward(
+        x, return_cache=True, training=False, update_stats=False)
+    signal = (torch.softmax(clean["logits"], dim=1)
+              - F.one_hot(y, 10).to(torch.float64))
+    edge = len(net.Q) - 2
+    observation = layerwise_causal_feedback_observation(
+        net, x, y, clean, signal, edge_index=edge, sigma=1e-3,
+        generator=torch.Generator().manual_seed(119))
+
+    recipient = observation["recipient"]
+    perturbations = [torch.zeros_like(value) for value in clean["hiddens"]]
+    probe = torch.zeros_like(clean["hiddens"][recipient], requires_grad=True)
+    perturbations[recipient] = probe
+    losses = F.cross_entropy(net.forward(
+        x, perturbations=perturbations, training=False,
+        update_stats=False)["logits"], y, reduction="none")
+    gradient = torch.autograd.grad(losses.sum(), probe)[0]
+    exact = (gradient * observation["direction"]).flatten(1).sum(1)
+    estimated = observation["directional"]
+    jvp_relative = float(
+        (estimated - exact).norm() / exact.norm().clamp_min(1e-30))
+    assert jvp_relative < 2e-3
+
+    left = CIFARHierarchicalFAResNet(
+        depth=8, base_width=2, seed=118, dtype=torch.float64,
+        normalization="batchnorm", residual_scale=1.0)
+    right = CIFARHierarchicalFAResNet(
+        depth=8, base_width=2, seed=118, dtype=torch.float64,
+        normalization="batchnorm", residual_scale=1.0)
+    before_q = [value.clone() for value in left.Q]
+    before_r = left.R_out.clone()
+    for weight in right.W:
+        weight.add_(torch.randn_like(weight))
+    right.W_out.add_(torch.randn_like(right.W_out))
+    left_metric = layerwise_causal_feedback_update(
+        left, observation, eta=0.1)
+    right_metric = layerwise_causal_feedback_update(
+        right, observation, eta=0.1)
+    locality_error = max(float((a - b).abs().max()) for a, b in zip(
+        left.Q + [left.R_out], right.Q + [right.R_out]))
+    assert locality_error < 1e-12
+    changed = [not torch.equal(a, b) for a, b in zip(before_q, left.Q)]
+    assert changed == [index == edge for index in range(len(left.Q))]
+    assert torch.equal(before_r, left.R_out)
+    assert all(math.isfinite(value) for value in left_metric.values()
+               if isinstance(value, float))
+    assert left_metric == right_metric
+
+    sweep_net = CIFARHierarchicalFAResNet(
+        depth=8, base_width=2, seed=120, dtype=torch.float64,
+        normalization="batchnorm", residual_scale=1.0)
+    sweep = layerwise_causal_bootstrap_sweep(
+        sweep_net, x[:2], y[:2], sigma=1e-3, eta=0.1,
+        generator=torch.Generator().manual_seed(121))
+    assert sweep["events"] == len(sweep_net.Q)
+    assert sweep["logical_batch_loss_queries"] == 2 * len(sweep_net.Q)
+    assert all(math.isfinite(value) for value in (
+        sweep["mean_field_prediction_target_cosine"],
+        sweep["mean_parameter_update_rms"],
+        sweep["max_parameter_update_rms"]))
+    return {
+        "layerwise_causal_per_example_jvp_relative_error": jvp_relative,
+        "layerwise_causal_update_forward_independence_error": locality_error,
+        "layerwise_causal_sweep_max_update_rms": (
+            sweep["max_parameter_update_rms"]),
+    }
+
+
 def normalized_response_mirror_checks():
     """Audit local response estimation and absence of W access in the update."""
     net = CIFARHierarchicalFAResNet(
@@ -1311,6 +1390,7 @@ def main():
     report.update(vectorizer_subspace_estimator_check())
     report.update(hierarchical_feedback_checks())
     report.update(hierarchical_parameter_calibration_checks())
+    report.update(layerwise_causal_bootstrap_checks())
     report.update(normalized_response_mirror_checks())
     report.update(kolen_pollack_checks())
     report.update(kp_mixed_traffic_checks())