baseline: add convolutional hierarchical feedback alignment

1 files changed, 157 insertions, 0 deletions

diff --git a/sdil/conv.py b/sdil/conv.py
index 384584e..f0592a6 100644
--- a/sdil/conv.py
+++ b/sdil/conv.py
@@ -411,6 +411,163 @@ class CIFARLocalResNet:
         return float(loss.detach())
 
 
+class CIFARHierarchicalFAResNet(CIFARLocalResNet):
+    """Residual-DAG feedback alignment with independent convolutional weights.
+
+    Feedback follows the actual child edges of the forward residual graph and
+    uses locally available ReLU/BatchNorm Jacobians, but every convolutional
+    feedback tensor is initialized independently and never reads its forward
+    counterpart.  This is a baseline and an infrastructure step for learned
+    hierarchical dendritic feedback, not an SDIL novelty claim.
+    """
+
+    def __init__(self, *args, feedback_seed=None, feedback_scale=1.0, **kwargs):
+        model_seed = kwargs.get("seed", 0)
+        super().__init__(*args, **kwargs)
+        if feedback_scale <= 0:
+            raise ValueError("feedback_scale must be positive")
+        generator = torch.Generator(device="cpu").manual_seed(
+            model_seed + 30011 if feedback_seed is None else feedback_seed)
+        self.Q = []
+        for weight in self.W:
+            fan_in = weight.shape[1] * weight.shape[2] * weight.shape[3]
+            value = (torch.randn(weight.shape, generator=generator)
+                     * (feedback_scale * math.sqrt(2.0 / fan_in)))
+            self.Q.append(value.to(device=weight.device, dtype=weight.dtype))
+        channels = self.W_out.shape[1]
+        self.R_out = (torch.randn(
+            channels, self.n_classes, generator=generator)
+            * (feedback_scale / math.sqrt(channels))).to(
+                device=self.W_out.device, dtype=self.W_out.dtype)
+
+    @property
+    def n_fixed_feedback_parameters(self):
+        # Q[0] maps the stem to pixels and is not used for hidden credit.
+        return (sum(value.numel() for value in self.Q[1:])
+                + self.R_out.numel())
+
+    @property
+    def apical_macs_per_example(self):
+        conv = 0
+        for weight, spec in zip(self.Q[1:], self.layer_specs[1:]):
+            out_channels, in_channels, kh, kw = weight.shape
+            _, height, width = spec.hidden_shape
+            conv += out_channels * height * width * in_channels * kh * kw
+        return int(conv + self.R_out.numel())
+
+    @staticmethod
+    def _option_a_shortcut_transpose(value, in_channels, stride, output_shape):
+        """Adjoint of the parameter-free option-A shortcut."""
+        out_channels = value.shape[1]
+        if in_channels > out_channels:
+            raise ValueError("option-A transpose cannot recover reduced channels")
+        missing = out_channels - in_channels
+        before = missing // 2
+        selected = value[:, before:before + in_channels]
+        if stride == 1:
+            if tuple(selected.shape) != tuple(output_shape):
+                raise ValueError("shortcut transpose shape mismatch")
+            return selected
+        result = value.new_zeros(output_shape)
+        result[:, :, ::2, ::2] = selected
+        return result
+
+    @torch.no_grad()
+    def hierarchical_teaching(self, output_signal, forward):
+        """Propagate teaching fields through independent feedback convolutions."""
+        caches = forward.get("caches")
+        hiddens = forward.get("hiddens")
+        if caches is None or hiddens is None:
+            raise ValueError("hierarchical feedback requires cached hidden states")
+        if len(hiddens) != self.n_hidden:
+            raise ValueError("hierarchical hidden population mismatch")
+        teaching = [torch.zeros_like(value) for value in hiddens]
+        spatial = hiddens[-1].shape[2] * hiddens[-1].shape[3]
+        teaching[-1].copy_(
+            (output_signal @ self.R_out.t())[:, :, None, None] / spatial)
+        for block in reversed(self.blocks):
+            first = block["first"]
+            second = block["second"]
+            parent = first - 1
+
+            second_gate = caches[second]["gate"].to(teaching[second].dtype)
+            second_delta = teaching[second] * second_gate
+            branch_delta, _, _ = self._normalization_backward(
+                second, second_delta * self.residual_scale,
+                caches[second]["normalization"])
+            teaching[first].add_(F.conv_transpose2d(
+                branch_delta, self.Q[second], stride=1, padding=1))
+            teaching[parent].add_(self._option_a_shortcut_transpose(
+                second_delta, block["in_channels"], block["stride"],
+                hiddens[parent].shape))
+
+            first_gate = caches[first]["gate"].to(teaching[first].dtype)
+            first_delta = teaching[first] * first_gate
+            first_delta, _, _ = self._normalization_backward(
+                first, first_delta, caches[first]["normalization"])
+            teaching[parent].add_(F.conv_transpose2d(
+                first_delta, self.Q[first], stride=block["stride"], padding=1,
+                output_padding=block["stride"] - 1))
+        return teaching
+
+
+def conv_hierarchical_step(net, x, y, config):
+    """One hierarchical-FA update using no reverse-mode graph or weight transport."""
+    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)
+        (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": False, "calibration": None,
+            "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]
+    for parameter in parameters:
+        parameter.requires_grad_(True)
+    forward = net.forward(x, return_cache=True, training=True, update_stats=False)
+    gradients = torch.autograd.grad(
+        F.cross_entropy(forward["logits"], y), forward["hiddens"])
+    batch = x.shape[0]
+    negative_gradients = [-batch * value.detach() for value in gradients]
+    with torch.no_grad():
+        output_error = (torch.softmax(forward["logits"], dim=1)
+                        - F.one_hot(y, net.n_classes).to(forward["logits"].dtype))
+        teaching = net.hierarchical_teaching(output_error, forward)
+        values = [float(F.cosine_similarity(
+            left.flatten(1), right.flatten(1), dim=1).mean())
+                  for left, right in zip(teaching, negative_gradients)]
+    for parameter in parameters:
+        parameter.requires_grad_(False)
+    return {
+        "normalization_state": "training_batch_stats_without_running_update",
+        "teaching_negative_gradient_cosine": values,
+        "raw_negative_gradient_cosine": values,
+        "innovation_negative_gradient_cosine": values,
+    }
+
+
 class CIFARSDILResNet(CIFARLocalResNet):
     """CIFAR local ResNet with per-unit apical vectorizers and predictors.