experiment: implement layerwise causal feedback bootstrap

2 files changed, 267 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())
diff --git a/sdil/conv.py b/sdil/conv.py
index 1034202..731cbab 100644
--- a/sdil/conv.py
+++ b/sdil/conv.py
@@ -1137,6 +1137,193 @@ def hierarchical_parameter_subspace_calibration(
 
 
 @torch.no_grad()
+def layerwise_causal_feedback_observation(
+        net, x, y, clean_forward, output_signal, edge_index=None,
+        sigma=1e-2, generator=None):
+    """Observe one hidden-credit target using per-example task perturbations.
+
+    ``edge_index=None`` denotes the dense readout feedback.  Otherwise only
+    the parent population receiving ``Q[edge_index]`` is perturbed.  Causal
+    queries use evaluation-mode normalization so each example's loss
+    difference depends only on its own perturbation.  The returned tensors are
+    a complete, detached local observation: the subsequent update need not and
+    does not read a forward parameter.
+    """
+    if not isinstance(net, CIFARHierarchicalFAResNet):
+        raise TypeError("layerwise causal calibration requires a hierarchical net")
+    if sigma <= 0:
+        raise ValueError("causal perturbation scale must be positive")
+    if generator is None:
+        generator = torch.Generator(device=x.device).manual_seed(0)
+    if len(clean_forward.get("hiddens", [])) != net.n_hidden:
+        raise ValueError("clean forward does not match hierarchical populations")
+    caches = clean_forward.get("caches")
+    if caches is None:
+        raise ValueError("layerwise causal calibration requires cached states")
+    if net.normalization == "batchnorm" and any(
+            cache["normalization"]["training"] for cache in caches):
+        raise ValueError(
+            "per-example causal observations require evaluation-mode BatchNorm")
+
+    teaching, contexts, recipients = net.hierarchical_teaching(
+        output_signal, clean_forward, return_edge_contexts=True)
+    perturbations = [torch.zeros_like(value)
+                     for value in clean_forward["hiddens"]]
+    batch = x.shape[0]
+    if edge_index is None:
+        recipient = net.n_hidden - 1
+        hidden = clean_forward["hiddens"][recipient]
+        random = torch.empty(
+            batch, hidden.shape[1], device=hidden.device,
+            dtype=hidden.dtype).bernoulli_(
+                0.5, generator=generator).mul_(2).sub_(1)
+        direction = random[:, :, None, None].expand_as(hidden)
+    else:
+        if not 1 <= edge_index < len(net.Q):
+            raise ValueError("feedback edge index is out of range")
+        recipient = recipients[edge_index]
+        hidden = clean_forward["hiddens"][recipient]
+        random = torch.empty_like(hidden).bernoulli_(
+            0.5, generator=generator).mul_(2).sub_(1)
+        direction = random
+    perturbations[recipient] = direction.mul(sigma)
+    plus = F.cross_entropy(net.forward(
+        x, perturbations=perturbations, training=False,
+        update_stats=False)["logits"], y, reduction="none")
+    perturbations[recipient] = direction.mul(-sigma)
+    minus = F.cross_entropy(net.forward(
+        x, perturbations=perturbations, training=False,
+        update_stats=False)["logits"], y, reduction="none")
+    directional = (plus - minus) / (2.0 * sigma)
+
+    if edge_index is None:
+        # A channel-constant hidden intervention is conjugate to the pooled
+        # feature vector.  The target therefore has the coefficient scale of
+        # output_signal @ R.T, before division by the number of spatial sites.
+        target = -directional[:, None] * random
+        prediction = output_signal @ net.R_out.t()
+        return {
+            "kind": "readout", "edge_index": None,
+            "recipient": recipient,
+            "target": target.clone(), "prediction": prediction.clone(),
+            "features": output_signal.clone(),
+            "direction": direction.clone(),
+            "directional": directional.clone(),
+        }
+
+    expand = directional.reshape(
+        batch, *([1] * (direction.ndim - 1)))
+    target = -expand * direction
+    return {
+        "kind": "convolution", "edge_index": int(edge_index),
+        "recipient": int(recipient), "target": target.clone(),
+        "prediction": teaching[recipient].clone(),
+        "context": contexts[edge_index].clone(),
+        "direction": direction.clone(),
+        "stride": int(net.layer_specs[edge_index].stride),
+        "padding": int(net.layer_specs[edge_index].padding),
+        "directional": directional.clone(),
+    }
+
+
+@torch.no_grad()
+def layerwise_causal_feedback_update(net, observation, eta=0.1, eps=1e-12):
+    """Apply one normalized local feedback update from a stored observation."""
+    if not isinstance(net, CIFARHierarchicalFAResNet):
+        raise TypeError("layerwise causal calibration requires a hierarchical net")
+    if not 0.0 < eta <= 1.0 or eps <= 0:
+        raise ValueError("invalid layerwise causal update hyperparameters")
+    target = observation["target"]
+    prediction = observation["prediction"]
+    error = target - prediction
+    kind = observation.get("kind")
+    if kind == "readout":
+        features = observation["features"]
+        normalizer = float(features.square().sum() / features.shape[0])
+        normalizer = max(normalizer, eps)
+        update = (error.t() @ features) / (features.shape[0] * normalizer)
+        if update.shape != net.R_out.shape:
+            raise ValueError("readout causal update shape mismatch")
+        net.R_out.add_(update, alpha=eta)
+        parameters = update.numel()
+        edge = "readout"
+    elif kind == "convolution":
+        index = int(observation["edge_index"])
+        context = observation["context"]
+        batch = error.shape[0]
+        recipient_spatial = error.shape[2] * error.shape[3]
+        normalizer = float(context.square().sum() / (batch * recipient_spatial))
+        normalizer = max(normalizer, eps)
+        update = torch.nn.grad.conv2d_weight(
+            error, net.Q[index].shape, context,
+            stride=int(observation["stride"]),
+            padding=int(observation["padding"]))
+        update.div_(batch * recipient_spatial * normalizer)
+        net.Q[index].add_(update, alpha=eta)
+        parameters = update.numel()
+        edge = index
+    else:
+        raise ValueError("unknown layerwise causal observation kind")
+
+    target_power = float(target.square().sum())
+    prediction_power = float(prediction.square().sum())
+    denominator = math.sqrt(target_power * prediction_power)
+    return {
+        "edge": edge,
+        "field_prediction_target_cosine": (
+            float((target * prediction).sum()) / denominator
+            if denominator else 0.0),
+        "target_rms": math.sqrt(target_power / target.numel()),
+        "prediction_rms": math.sqrt(prediction_power / prediction.numel()),
+        "field_error_rms": math.sqrt(float(error.square().mean())),
+        "parameter_update_rms": math.sqrt(
+            float(update.square().sum()) / parameters),
+        "local_context_power": normalizer,
+    }
+
+
+@torch.no_grad()
+def layerwise_causal_bootstrap_sweep(
+        net, x, y, sigma=1e-2, eta=0.1, generator=None):
+    """Calibrate readout then convolutional feedback in reverse DAG order."""
+    if generator is None:
+        generator = torch.Generator(device=x.device).manual_seed(0)
+    clean = net.forward(
+        x, return_cache=True, training=False, update_stats=False)
+    output_signal = (torch.softmax(clean["logits"], dim=1)
+                     - F.one_hot(y, net.n_classes).to(clean["logits"].dtype))
+    metrics = []
+    observation = layerwise_causal_feedback_observation(
+        net, x, y, clean, output_signal, edge_index=None,
+        sigma=sigma, generator=generator)
+    metrics.append(layerwise_causal_feedback_update(
+        net, observation, eta=eta))
+    for index in reversed(range(1, len(net.Q))):
+        # The forward cache is unchanged by feedback learning.  The teaching
+        # pass inside the observation is recomputed after every Q update, so
+        # the newly calibrated child field bootstraps its immediate parent.
+        observation = layerwise_causal_feedback_observation(
+            net, x, y, clean, output_signal, edge_index=index,
+            sigma=sigma, generator=generator)
+        metrics.append(layerwise_causal_feedback_update(
+            net, observation, eta=eta))
+    return {
+        "events": len(metrics),
+        "logical_batch_loss_queries": 2 * len(metrics),
+        "per_example_causal_observations": x.shape[0] * len(metrics),
+        "mean_field_prediction_target_cosine": sum(
+            value["field_prediction_target_cosine"] for value in metrics)
+            / len(metrics),
+        "mean_parameter_update_rms": sum(
+            value["parameter_update_rms"] for value in metrics)
+            / len(metrics),
+        "max_parameter_update_rms": max(
+            value["parameter_update_rms"] for value in metrics),
+        "edges": metrics,
+    }
+
+
+@torch.no_grad()
 def hierarchical_mirror_observations(net, batch_size=1, noise_std=1.0,
                                      generator=None):
     """Generate local bias-blocked probe/child-response pairs.