experiment: implement layerwise causal feedback bootstrap

1 files changed, 187 insertions, 0 deletions

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.