path: root/sdil/conv_crossover.py

"""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 CIFARPEPITAResNet(CIFARLocalResNet):
    """Two-presentation PEPITA on the matched residual forward topology."""

    def __init__(self, *args, projection_scale=0.05,
                 projection_seed=1731, **kwargs):
        super().__init__(*args, **kwargs)
        if projection_scale <= 0:
            raise ValueError("PEPITA projection scale must be positive")
        generator = torch.Generator(device="cpu").manual_seed(projection_seed)
        input_units = 3 * 32 * 32
        limit = (6.0 / input_units) ** 0.5 * projection_scale
        projection = (
            2.0 * torch.rand(
                self.n_classes, 3, 32, 32, generator=generator) - 1.0
        ) * limit
        self.input_feedback = projection.to(
            device=self.device, dtype=self.dtype)

    @property
    def n_fixed_feedback_parameters(self):
        return self.input_feedback.numel()

    @torch.no_grad()
    def pepita_ascent_directions(self, clean, modulated, modulated_error):
        """Return the explicit first-minus-second PEPITA correlations.

        The post-activation difference directly multiplies the modulated
        presynaptic activity, as in PEPITA/ERIN; it is not differentiated
        through the ReLU.  BatchNorm's current-layer Jacobian and a residual
        branch's fixed multiplier remain part of that local synaptic
        eligibility.  Convolutional correlations follow the reference code's
        additional average over spatial positions.
        """
        batch = modulated_error.shape[0]
        directions = []
        gamma_directions = []
        beta_directions = []
        for index, (clean_hidden, modulated_hidden, cache, weight, spec) in (
                enumerate(zip(
                    clean["hiddens"], modulated["hiddens"],
                    modulated["caches"], self.W, self.layer_specs))):
            field = clean_hidden - modulated_hidden
            local_field = field * spec.branch_scale
            local_field, gamma_direction, beta_direction = (
                self._normalization_backward(
                    index, local_field, cache["normalization"]))
            spatial = local_field.shape[2] * local_field.shape[3]
            correlation = torch.nn.grad.conv2d_weight(
                cache["pre"], weight.shape, local_field,
                stride=spec.stride, padding=spec.padding)
            directions.append(-correlation / (batch * spatial))
            if gamma_direction is not None:
                gamma_directions.append(
                    -gamma_direction / (batch * spatial))
                beta_directions.append(
                    -beta_direction / (batch * spatial))
        output_weight = -(
            modulated_error.t() @ modulated["features"]) / batch
        output_bias = -modulated_error.mean(dim=0)
        return (
            directions, gamma_directions, beta_directions,
            output_weight, output_bias)

    def pepita_step(self, image, labels, eta, eta_output=None,
                    momentum=0.0, weight_decay=0.0):
        """One architecture-compatible PEPITA/ERIN local update."""
        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)
            clean_error = torch.softmax(clean["logits"], dim=1) - one_hot
            input_error = torch.einsum(
                "bc,cijk->bijk", clean_error, self.input_feedback)
            modulated = self.forward(
                image + input_error, return_cache=True,
                training=True, update_stats=False)
            modulated_error = (
                torch.softmax(modulated["logits"], dim=1) - one_hot)
            (directions, gamma_directions, beta_directions,
             output_weight, output_bias) = self.pepita_ascent_directions(
                 clean, modulated, modulated_error)
            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)
            loss = F.cross_entropy(clean["logits"], labels)
        return float(loss)


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)