baseline: add matched ResNet Dual Propagation

1 files changed, 214 insertions, 0 deletions

diff --git a/sdil/conv_crossover.py b/sdil/conv_crossover.py
new file mode 100644
index 0000000..d455f21
--- /dev/null
+++ b/sdil/conv_crossover.py
@@ -0,0 +1,214 @@
+"""Matched non-backprop adapters for the audited CIFAR ResNet topology.
+
+These adapters reuse :class:`CIFARLocalResNet`'s forward tensors, BatchNorm,
+option-A shortcuts, and local optimizer.  They are kept separate from the
+established SDIL/KP implementation so a crossover baseline cannot silently
+change the already confirmed forward model.
+"""
+import torch
+import torch.nn.functional as F
+
+from .conv import CIFARLocalResNet
+
+
+class CIFARDualPropResNet(CIFARLocalResNet):
+    """Dual Propagation on the residual DAG with author DP-transpose updates.
+
+    ``s_plus`` and ``s_minus`` are initialized to the ordinary forward states.
+    A ``fwK`` pass updates every residual-DAG node in topological order.  The
+    feedforward drive uses the forward edge, while the difference of each
+    child state is transported through the exact transpose of that same edge.
+    This is intentional symmetric feedback in the Dual Propagation baseline,
+    not a claim of weight-transport-free learning.
+    """
+
+    def __init__(self, *args, alpha=0.0, dp_beta=0.1,
+                 inference_passes=16, **kwargs):
+        super().__init__(*args, **kwargs)
+        if not 0.0 <= alpha <= 1.0:
+            raise ValueError("Dual Propagation alpha must lie in [0, 1]")
+        if dp_beta <= 0 or inference_passes < 1:
+            raise ValueError("invalid Dual Propagation inference settings")
+        self.dp_alpha = float(alpha)
+        self.dp_beta = float(dp_beta)
+        self.dp_inference_passes = int(inference_passes)
+        self._node_kind = {0: ("stem",)}
+        self._outgoing = {index: [] for index in range(self.n_hidden)}
+        for block in self.blocks:
+            first = block["first"]
+            second = block["second"]
+            parent = first - 1
+            self._node_kind[first] = ("first", parent)
+            self._node_kind[second] = (
+                "second", first, parent, block["out_channels"],
+                block["stride"])
+            self._outgoing[parent].append((first, "conv"))
+            self._outgoing[first].append((second, "conv"))
+            self._outgoing[parent].append((second, "shortcut"))
+
+    @staticmethod
+    def _option_a_shortcut_transpose(value, input_shape, stride):
+        """Adjoint of the parameter-free option-A shortcut."""
+        in_channels = input_shape[1]
+        out_channels = value.shape[1]
+        missing = out_channels - in_channels
+        if missing < 0:
+            raise ValueError("option-A transpose cannot increase input width")
+        before = missing // 2
+        selected = value[:, before:before + in_channels]
+        if stride == 1:
+            if tuple(selected.shape) != tuple(input_shape):
+                raise ValueError("option-A transpose shape mismatch")
+            return selected
+        result = value.new_zeros(input_shape)
+        result[:, :, ::2, ::2] = selected
+        return result
+
+    def _node_prediction(self, index, states, image):
+        """Return the unrectified local prediction and its edge cache."""
+        kind = self._node_kind[index]
+        if kind[0] == "stem":
+            pre = image
+            shortcut = None
+        elif kind[0] == "first":
+            pre = states[kind[1]]
+            shortcut = None
+        else:
+            pre = states[kind[1]]
+            shortcut = self._option_a_shortcut(
+                states[kind[2]], kind[3], kind[4])
+        spec = self.layer_specs[index]
+        convolution = F.conv2d(
+            pre, self.W[index], stride=spec.stride, padding=spec.padding)
+        normalized, normalization = self._normalize(
+            index, convolution, training=True, update_stats=False)
+        prediction = (
+            shortcut + spec.branch_scale * normalized
+            if shortcut is not None else normalized)
+        return prediction, {
+            "pre": pre,
+            "normalization": normalization,
+            "shortcut": shortcut,
+        }
+
+    def _edge_transpose(self, child, edge_kind, field, states, image):
+        """Apply one residual-DAG edge transpose to a child state field."""
+        kind = self._node_kind[child]
+        if edge_kind == "shortcut":
+            if kind[0] != "second":
+                raise AssertionError("only a second convolution has a shortcut")
+            return self._option_a_shortcut_transpose(
+                field, states[kind[2]].shape, kind[4])
+        _, cache = self._node_prediction(child, states, image)
+        spec = self.layer_specs[child]
+        local_field = field * spec.branch_scale
+        local_field, _, _ = self._normalization_backward(
+            child, local_field, cache["normalization"])
+        return torch.nn.grad.conv2d_input(
+            cache["pre"].shape, self.W[child], local_field,
+            stride=spec.stride, padding=spec.padding)
+
+    def _outgoing_feedback(self, index, deltas, states, image):
+        result = torch.zeros_like(states[index])
+        for child, edge_kind in self._outgoing[index]:
+            result.add_(self._edge_transpose(
+                child, edge_kind, deltas[child], states, image))
+        if index == self.n_hidden - 1:
+            spatial = states[index].shape[2] * states[index].shape[3]
+            result.add_(
+                (deltas[-1] @ self.W_out)[:, :, None, None] / spatial)
+        return result
+
+    def infer_dual_states(self, image, one_hot, clean_forward=None):
+        """Run the author ``fwK`` DP-transpose state updates on the DAG."""
+        if clean_forward is None:
+            clean_forward = self.forward(
+                image, return_cache=True, training=True, update_stats=False)
+        plus = [
+            value.detach().clone() for value in clean_forward["hiddens"]]
+        minus = [value.detach().clone() for value in plus]
+        plus.append(clean_forward["logits"].detach().clone())
+        minus.append(clean_forward["logits"].detach().clone())
+        fixed_prediction = clean_forward["logits"].detach()
+        alpha = self.dp_alpha
+        for _ in range(self.dp_inference_passes):
+            for index in range(self.n_hidden):
+                states = [
+                    alpha * positive + (1.0 - alpha) * negative
+                    for positive, negative in zip(plus[:-1], minus[:-1])]
+                prediction, _ = self._node_prediction(index, states, image)
+                deltas = [
+                    positive - negative
+                    for positive, negative in zip(plus, minus)]
+                feedback = self._outgoing_feedback(
+                    index, deltas, states, image)
+                plus[index] = F.relu(
+                    prediction + (1.0 - alpha) * feedback)
+                minus[index] = F.relu(prediction - alpha * feedback)
+            states = [
+                alpha * positive + (1.0 - alpha) * negative
+                for positive, negative in zip(plus[:-1], minus[:-1])]
+            features = states[-1].mean(dim=(2, 3))
+            prediction = features @ self.W_out.t() + self.b_out
+            output_field = self.dp_beta * (
+                torch.softmax(fixed_prediction, dim=1) - one_hot)
+            plus[-1] = prediction - (1.0 - alpha) * output_field
+            minus[-1] = prediction + alpha * output_field
+        return plus, minus
+
+    @torch.no_grad()
+    def dualprop_ascent_directions(self, image, plus, minus):
+        """Evaluate the local DP contrastive correlations without autograd."""
+        alpha = self.dp_alpha
+        beta = self.dp_beta
+        states = [
+            alpha * positive + (1.0 - alpha) * negative
+            for positive, negative in zip(plus[:-1], minus[:-1])]
+        deltas = [
+            (positive - negative) / beta
+            for positive, negative in zip(plus, minus)]
+        directions = []
+        gamma_directions = []
+        beta_directions = []
+        batch = image.shape[0]
+        for index in range(self.n_hidden):
+            _, cache = self._node_prediction(index, states, image)
+            spec = self.layer_specs[index]
+            local_field = deltas[index] * spec.branch_scale
+            local_field, gamma_direction, beta_direction = (
+                self._normalization_backward(
+                    index, local_field, cache["normalization"]))
+            direction = torch.nn.grad.conv2d_weight(
+                cache["pre"], self.W[index].shape, local_field,
+                stride=spec.stride, padding=spec.padding)
+            directions.append(direction / batch)
+            if gamma_direction is not None:
+                gamma_directions.append(gamma_direction / batch)
+                beta_directions.append(beta_direction / batch)
+        features = states[-1].mean(dim=(2, 3))
+        output_weight = deltas[-1].t() @ features / batch
+        output_bias = deltas[-1].mean(dim=0)
+        return (
+            directions, gamma_directions, beta_directions,
+            output_weight, output_bias)
+
+    def dualprop_step(self, image, labels, eta, eta_output=None,
+                      momentum=0.0, weight_decay=0.0):
+        """One fully local DP-transpose update on a matched ResNet batch."""
+        one_hot = F.one_hot(labels, self.n_classes).to(image.dtype)
+        with torch.no_grad():
+            clean = self.forward(
+                image, return_cache=True, training=True, update_stats=True)
+            loss = F.cross_entropy(clean["logits"], labels)
+            plus, minus = self.infer_dual_states(
+                image, one_hot, clean_forward=clean)
+            (directions, gamma_directions, beta_directions,
+             output_weight, output_bias) = self.dualprop_ascent_directions(
+                 image, plus, minus)
+            self.apply_ascent(
+                directions, output_weight, output_bias, eta,
+                eta_output=eta_output, momentum=momentum,
+                weight_decay=weight_decay,
+                gamma_directions=gamma_directions,
+                beta_directions=beta_directions)
+        return float(loss)