algorithm: calibrate hierarchical feedback parameter subspaces

1 files changed, 192 insertions, 3 deletions

diff --git a/sdil/conv.py b/sdil/conv.py
index f0592a6..55cdc62 100644
--- a/sdil/conv.py
+++ b/sdil/conv.py
@@ -473,8 +473,16 @@ class CIFARHierarchicalFAResNet(CIFARLocalResNet):
         return result
 
     @torch.no_grad()
-    def hierarchical_teaching(self, output_signal, forward):
-        """Propagate teaching fields through independent feedback convolutions."""
+    def hierarchical_teaching(self, output_signal, forward,
+                              return_edge_contexts=False):
+        """Propagate teaching fields through independent feedback convolutions.
+
+        When requested, ``edge_contexts[l]`` is the local child field consumed
+        by ``Q[l]`` and ``recipients[l]`` is the hidden population receiving
+        that transposed-convolution contribution.  These values expose the
+        sufficient statistics for causal feedback calibration without reading
+        any forward tensor.
+        """
         caches = forward.get("caches")
         hiddens = forward.get("hiddens")
         if caches is None or hiddens is None:
@@ -482,6 +490,8 @@ class CIFARHierarchicalFAResNet(CIFARLocalResNet):
         if len(hiddens) != self.n_hidden:
             raise ValueError("hierarchical hidden population mismatch")
         teaching = [torch.zeros_like(value) for value in hiddens]
+        edge_contexts = [None for _ in self.Q]
+        recipients = [None for _ in self.Q]
         spatial = hiddens[-1].shape[2] * hiddens[-1].shape[3]
         teaching[-1].copy_(
             (output_signal @ self.R_out.t())[:, :, None, None] / spatial)
@@ -495,6 +505,8 @@ class CIFARHierarchicalFAResNet(CIFARLocalResNet):
             branch_delta, _, _ = self._normalization_backward(
                 second, second_delta * self.residual_scale,
                 caches[second]["normalization"])
+            edge_contexts[second] = branch_delta
+            recipients[second] = first
             teaching[first].add_(F.conv_transpose2d(
                 branch_delta, self.Q[second], stride=1, padding=1))
             teaching[parent].add_(self._option_a_shortcut_transpose(
@@ -505,12 +517,150 @@ class CIFARHierarchicalFAResNet(CIFARLocalResNet):
             first_delta = teaching[first] * first_gate
             first_delta, _, _ = self._normalization_backward(
                 first, first_delta, caches[first]["normalization"])
+            edge_contexts[first] = first_delta
+            recipients[first] = parent
             teaching[parent].add_(F.conv_transpose2d(
                 first_delta, self.Q[first], stride=block["stride"], padding=1,
                 output_padding=block["stride"] - 1))
+        if return_edge_contexts:
+            if any(value is None for value in edge_contexts[1:]):
+                raise AssertionError("hierarchical edge context is incomplete")
+            return teaching, edge_contexts, recipients
         return teaching
 
 
+@torch.no_grad()
+def hierarchical_parameter_subspace_calibration(
+        net, x, y, clean_forward, output_signal, sigma=1e-2,
+        n_directions=1, eta=1e-3, generator=None, return_diagnostics=False):
+    """Calibrate the residual-DAG feedback maps with two causal queries.
+
+    A Rademacher tensor is drawn in every Q/R parameter space.  Each tensor is
+    applied only to its locally available child field, producing one candidate
+    intervention at the edge's parent population.  All independent candidate
+    fields are injected in the same antithetic pair.  Multiplying the scalar
+    loss derivative back into each random tensor gives an unbiased estimate of
+    that edge's causal target moment; subtracting the current predicted moment
+    is the exact local squared-field delta rule.  No forward weight or reverse
+    differentiation is used.
+    """
+    if not isinstance(net, CIFARHierarchicalFAResNet):
+        raise TypeError("hierarchical calibration requires a hierarchical net")
+    if sigma <= 0 or n_directions < 1 or eta < 0:
+        raise ValueError("invalid hierarchical calibration hyperparameters")
+    if generator is None:
+        generator = torch.Generator(device=x.device).manual_seed(0)
+    hiddens = clean_forward.get("hiddens")
+    if hiddens is None or len(hiddens) != net.n_hidden:
+        raise ValueError("clean forward does not match hierarchical populations")
+    teaching, contexts, recipients = net.hierarchical_teaching(
+        output_signal, clean_forward, return_edge_contexts=True)
+    batch = x.shape[0]
+    target_q = [torch.zeros_like(value) if index else None
+                for index, value in enumerate(net.Q)]
+    target_r = torch.zeros_like(net.R_out)
+    diagnostic_directions = []
+    diagnostic_derivatives = []
+    for _ in range(n_directions):
+        random_q = [None]
+        random_q.extend([
+            torch.empty_like(value).bernoulli_(
+                0.5, generator=generator).mul_(2).sub_(1)
+            for value in net.Q[1:]
+        ])
+        random_r = torch.empty_like(net.R_out).bernoulli_(
+            0.5, generator=generator).mul_(2).sub_(1)
+        interventions = [torch.zeros_like(value) for value in hiddens]
+        spatial_out = hiddens[-1].shape[2] * hiddens[-1].shape[3]
+        interventions[-1].add_(
+            (output_signal @ random_r.t())[:, :, None, None] / spatial_out)
+        for index in range(1, len(net.Q)):
+            spec = net.layer_specs[index]
+            recipient = recipients[index]
+            contribution = F.conv_transpose2d(
+                contexts[index], random_q[index], stride=spec.stride,
+                padding=spec.padding, output_padding=spec.stride - 1)
+            if contribution.shape != interventions[recipient].shape:
+                raise AssertionError("hierarchical intervention shape mismatch")
+            interventions[recipient].add_(contribution)
+        plus = F.cross_entropy(net.forward(
+            x, perturbations=[sigma * value for value in interventions],
+            training=True, update_stats=False)["logits"], y)
+        minus = F.cross_entropy(net.forward(
+            x, perturbations=[-sigma * value for value in interventions],
+            training=True, update_stats=False)["logits"], y)
+        # The parameter maps are batch-shared.  Scale the mean-loss derivative
+        # into a summed-loss hidden signal, then average its moment over B and
+        # recipient spatial sites to keep one eta meaningful across stages.
+        directional = (plus - minus) * batch / (2.0 * sigma)
+        target_r.add_(random_r, alpha=-float(directional) / (
+            batch * n_directions))
+        for index in range(1, len(net.Q)):
+            recipient = recipients[index]
+            spatial = (hiddens[recipient].shape[2]
+                       * hiddens[recipient].shape[3])
+            target_q[index].add_(
+                random_q[index], alpha=-float(directional) / (
+                    batch * spatial * n_directions))
+        if return_diagnostics:
+            diagnostic_directions.append({
+                "hidden": interventions, "Q": random_q, "R": random_r})
+            diagnostic_derivatives.append({
+                "scaled_directional": directional,
+                "coupling": "summed_batch_objective",
+            })
+
+    prediction_q = [None]
+    for index in range(1, len(net.Q)):
+        recipient = recipients[index]
+        spec = net.layer_specs[index]
+        spatial = (hiddens[recipient].shape[2]
+                   * hiddens[recipient].shape[3])
+        prediction_q.append(torch.nn.grad.conv2d_weight(
+            teaching[recipient], net.Q[index].shape, contexts[index],
+            stride=spec.stride, padding=spec.padding) / (batch * spatial))
+    prediction_r = (teaching[-1].mean(dim=(2, 3)).t()
+                    @ output_signal) / batch
+    errors_q = [None]
+    errors_q.extend([
+        target_q[index] - prediction_q[index]
+        for index in range(1, len(net.Q))
+    ])
+    error_r = target_r - prediction_r
+    for index in range(1, len(net.Q)):
+        net.Q[index].add_(errors_q[index], alpha=eta)
+    net.R_out.add_(error_r, alpha=eta)
+
+    target_power = float(target_r.square().sum())
+    prediction_power = float(prediction_r.square().sum())
+    error_power = float(error_r.square().sum())
+    dot = float((target_r * prediction_r).sum())
+    parameters = target_r.numel()
+    for index in range(1, len(net.Q)):
+        target_power += float(target_q[index].square().sum())
+        prediction_power += float(prediction_q[index].square().sum())
+        error_power += float(errors_q[index].square().sum())
+        dot += float((target_q[index] * prediction_q[index]).sum())
+        parameters += target_q[index].numel()
+    denominator = math.sqrt(target_power * prediction_power)
+    calibration = {
+        "calibration_mse": error_power / parameters,
+        "target_power": target_power / parameters,
+        "prediction_target_cosine": dot / denominator if denominator else 0.0,
+        "parameter_update_rms": math.sqrt(error_power / parameters),
+    }
+    if return_diagnostics:
+        return calibration, {
+            "directions": diagnostic_directions,
+            "directional_derivatives": diagnostic_derivatives,
+            "teaching": teaching, "contexts": contexts,
+            "recipients": recipients, "target_Q": target_q,
+            "prediction_Q": prediction_q, "target_R": target_r,
+            "prediction_R": prediction_r,
+        }
+    return calibration
+
+
 def conv_hierarchical_step(net, x, y, config):
     """One hierarchical-FA update using no reverse-mode graph or weight transport."""
     config.validate()
@@ -541,6 +691,44 @@ def conv_hierarchical_step(net, x, y, config):
         }
 
 
+def conv_learned_hierarchical_step(net, x, y, config, step, generator=None):
+    """One task update with optional causal calibration of hierarchical Q/R."""
+    config.validate()
+    with torch.no_grad():
+        forward = net.forward(
+            x, return_cache=True, training=True, update_stats=True)
+        logits = forward["logits"]
+        loss = F.cross_entropy(logits, y)
+        output_error = (torch.softmax(logits, dim=1)
+                        - F.one_hot(y, net.n_classes).to(logits.dtype))
+        teaching = net.hierarchical_teaching(output_error, forward)
+        total_units = sum(value.numel() for value in teaching)
+        teaching_rms = math.sqrt(
+            sum(float(value.square().sum()) for value in teaching) / total_units)
+        did_perturb = config.learn_A and step % config.pert_every == 0
+        calibration = None
+        if did_perturb:
+            calibration = hierarchical_parameter_subspace_calibration(
+                net, x, y, forward, output_error, sigma=config.pert_sigma,
+                n_directions=config.pert_directions, eta=config.eta_A,
+                generator=generator)
+        (directions, gamma_directions, beta_directions,
+         output_weight, output_bias) = net.local_ascent_directions(
+            teaching, output_error, forward)
+        net.apply_ascent(
+            directions, output_weight, output_bias,
+            eta_hidden=config.eta, eta_output=config.eta_output,
+            momentum=config.momentum, weight_decay=config.weight_decay,
+            gamma_directions=gamma_directions,
+            beta_directions=beta_directions)
+        return {
+            "loss": float(loss), "did_perturb": did_perturb,
+            "calibration": calibration, "predictor_mse": None,
+            "teaching_rms": teaching_rms, "raw_apical_rms": teaching_rms,
+            "innovation_rms": teaching_rms,
+        }
+
+
 def conv_hierarchical_alignment_report(net, x, y):
     """Audit hierarchical teaching against exact hidden gradients."""
     parameters = net.W + net.gamma + net.beta + [net.W_out, net.b_out]
@@ -1081,7 +1269,8 @@ class ConvSDILConfig:
         if self.pert_every < 1 or self.pert_directions < 1:
             raise ValueError("perturbation cadence/directions must be positive")
         if self.apical_calibration_mode not in (
-                "unit_targets", "channel_subspace", "vectorizer_subspace"):
+                "unit_targets", "channel_subspace", "vectorizer_subspace",
+                "hierarchical_parameter_subspace"):
             raise ValueError("unknown apical calibration mode")
         if self.direct_node_perturbation and self.pert_every != 1:
             raise ValueError("direct node perturbation requires a target every step")