baseline: add matched ResNet Dual Propagation

2 files changed, 308 insertions, 0 deletions

diff --git a/experiments/resnet_crossover_smoke.py b/experiments/resnet_crossover_smoke.py
new file mode 100644
index 0000000..fdc92fa
--- /dev/null
+++ b/experiments/resnet_crossover_smoke.py
@@ -0,0 +1,94 @@
+#!/usr/bin/env python3
+"""Deterministic equation audits for matched ResNet crossover adapters."""
+import json
+import os
+import sys
+
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from sdil.conv import CIFARLocalResNet
+from sdil.conv_crossover import CIFARDualPropResNet
+
+
+def relative_error(actual, expected):
+    numerator = torch.linalg.vector_norm(actual - expected)
+    denominator = torch.linalg.vector_norm(expected).clamp_min(1e-30)
+    return float(numerator / denominator)
+
+
+def audit_dualprop():
+    common = dict(
+        depth=8, base_width=2, seed=71, normalization="batchnorm",
+        residual_scale=1.0, dtype=torch.float64)
+    reference = CIFARLocalResNet(**common)
+    net = CIFARDualPropResNet(
+        **common, alpha=0.0, dp_beta=0.1, inference_passes=2)
+    generator = torch.Generator().manual_seed(72)
+    image = torch.randn(3, 3, 32, 32, generator=generator,
+                        dtype=torch.float64)
+    labels = torch.tensor([0, 3, 7])
+    one_hot = torch.nn.functional.one_hot(labels, 10).to(torch.float64)
+    with torch.no_grad():
+        reference_output = reference.forward(
+            image, training=True, update_stats=False)["logits"]
+        net_output = net.forward(
+            image, training=True, update_stats=False)["logits"]
+        clean = net.forward(
+            image, return_cache=True, training=True, update_stats=False)
+        plus, minus = net.infer_dual_states(image, one_hot, clean)
+    forward_error = float(torch.max(torch.abs(
+        reference_output - net_output)))
+
+    parameters = (
+        net.W + net.gamma + net.beta + [net.W_out, net.b_out])
+    for parameter in parameters:
+        parameter.requires_grad_(True)
+    alpha = net.dp_alpha
+    beta = net.dp_beta
+    states = [
+        (alpha * positive + (1.0 - alpha) * negative).detach()
+        for positive, negative in zip(plus[:-1], minus[:-1])]
+    deltas = [
+        ((positive - negative) / beta).detach()
+        for positive, negative in zip(plus, minus)]
+    objective = image.new_zeros(())
+    for index in range(net.n_hidden):
+        prediction, _ = net._node_prediction(index, states, image)
+        objective -= torch.sum(deltas[index] * prediction) / image.shape[0]
+    features = states[-1].mean(dim=(2, 3))
+    output_prediction = features @ net.W_out.t() + net.b_out
+    objective -= torch.sum(deltas[-1] * output_prediction) / image.shape[0]
+    gradients = torch.autograd.grad(objective, parameters)
+    (directions, gamma_directions, beta_directions,
+     output_weight, output_bias) = net.dualprop_ascent_directions(
+         image, plus, minus)
+    actual = (
+        directions + gamma_directions + beta_directions
+        + [output_weight, output_bias])
+    errors = [
+        relative_error(direction, -gradient)
+        for direction, gradient in zip(actual, gradients)]
+    for parameter in parameters:
+        parameter.requires_grad_(False)
+    if forward_error >= 1e-12 or max(errors) >= 2e-12:
+        raise AssertionError({
+            "forward_error": forward_error,
+            "direction_errors": errors,
+        })
+    return {
+        "matched_forward_max_absolute_error": forward_error,
+        "contrastive_direction_max_relative_error": max(errors),
+        "num_audited_parameter_tensors": len(errors),
+        "uses_symmetric_forward_edge_transposes": True,
+        "uses_reverse_task_loss_graph": False,
+    }
+
+
+def main():
+    print(json.dumps({"dualprop": audit_dualprop()},
+                     indent=2, sort_keys=True))
+
+
+if __name__ == "__main__":
+    main()
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)