baseline: add matched ResNet Forward-Forward

2 files changed, 259 insertions, 1 deletions

diff --git a/experiments/resnet_crossover_smoke.py b/experiments/resnet_crossover_smoke.py
index aadede5..8cdcb0c 100644
--- a/experiments/resnet_crossover_smoke.py
+++ b/experiments/resnet_crossover_smoke.py
@@ -8,7 +8,11 @@ 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, CIFARPEPITAResNet
+from sdil.conv_crossover import (
+    CIFARDualPropResNet,
+    CIFARForwardForwardResNet,
+    CIFARPEPITAResNet,
+)
 
 
 def relative_error(actual, expected):
@@ -167,9 +171,88 @@ def audit_pepita():
     }
 
 
+def audit_forward_forward():
+    common = dict(
+        depth=8, base_width=2, seed=91, normalization="batchnorm",
+        residual_scale=1.0, dtype=torch.float64)
+    reference = CIFARLocalResNet(**common)
+    net = CIFARForwardForwardResNet(
+        **common, threshold=2.0, learning_rate=0.03,
+        score_from_layer=1)
+    if net.n_forward_parameters != reference.n_forward_parameters:
+        raise AssertionError("Forward-Forward changed forward parameter count")
+    generator = torch.Generator().manual_seed(92)
+    image = torch.randn(4, 3, 32, 32, generator=generator,
+                        dtype=torch.float64)
+    labels = torch.tensor([1, 2, 6, 9])
+    negative_labels = (labels + 3) % 10
+    target = 1
+    with torch.no_grad():
+        positive = net.ff_forward(
+            net.ff_overlay(image, labels),
+            training=True, update_stats=False)
+        negative = net.ff_forward(
+            net.ff_overlay(image, negative_labels),
+            training=True, update_stats=False)
+    positive_output, negative_output = net._ff_local_outputs(
+        target, positive, negative)
+    axes = tuple(range(1, positive_output.ndim))
+    positive_goodness = positive_output.square().mean(dim=axes)
+    negative_goodness = negative_output.square().mean(dim=axes)
+    explicit_loss = (
+        torch.nn.functional.softplus(-positive_goodness + net.ff_threshold)
+        + torch.nn.functional.softplus(
+            negative_goodness - net.ff_threshold)
+    ).mean()
+    parameters = (
+        net.W + net.gamma + net.beta + [net.W_out, net.b_out])
+    before = [parameter.detach().clone() for parameter in parameters]
+    metrics = net.ff_train_layer(
+        target, image, labels, learning_rate=0.03,
+        negative_labels=negative_labels)
+    changed = [
+        not torch.equal(old, parameter.detach())
+        for old, parameter in zip(before, parameters)]
+    target_indices = {target}
+    if net.normalization == "batchnorm":
+        target_indices.update({
+            len(net.W) + target,
+            len(net.W) + len(net.gamma) + target,
+        })
+    non_target_changes = [
+        index for index, value in enumerate(changed)
+        if value and index not in target_indices]
+    non_target_gradients = [
+        index for index, parameter in enumerate(parameters)
+        if index not in target_indices and parameter.grad is not None]
+    scores = net.ff_candidate_scores(image)
+    loss_error = abs(metrics["loss"] - float(explicit_loss.detach()))
+    if (loss_error >= 1e-12 or non_target_changes
+            or non_target_gradients or not any(
+                changed[index] for index in target_indices)
+            or scores.shape != (4, 10)
+            or not torch.isfinite(scores).all()):
+        raise AssertionError({
+            "loss_error": loss_error,
+            "non_target_changes": non_target_changes,
+            "non_target_gradients": non_target_gradients,
+            "changed": changed,
+            "score_shape": list(scores.shape),
+        })
+    return {
+        "local_objective_absolute_error": loss_error,
+        "non_target_parameter_changes": len(non_target_changes),
+        "non_target_gradients": len(non_target_gradients),
+        "candidate_score_shape": list(scores.shape),
+        "matched_forward_parameter_count": net.n_forward_parameters,
+        "uses_only_target_layer_autograd": True,
+    }
+
+
 def main():
     print(json.dumps({
         "dualprop": audit_dualprop(),
+        "forward_forward": audit_forward_forward(),
         "pepita": audit_pepita(),
     },
                      indent=2, sort_keys=True))
diff --git a/sdil/conv_crossover.py b/sdil/conv_crossover.py
index 329ad3b..092bffe 100644
--- a/sdil/conv_crossover.py
+++ b/sdil/conv_crossover.py
@@ -102,6 +102,181 @@ class CIFARPEPITAResNet(CIFARLocalResNet):
         return float(loss)
 
 
+class CIFARForwardForwardResNet(CIFARLocalResNet):
+    """Greedy supervised Forward-Forward on the residual parameter topology."""
+
+    def __init__(self, *args, threshold=2.0, learning_rate=0.03,
+                 score_from_layer=1, **kwargs):
+        super().__init__(*args, **kwargs)
+        if learning_rate <= 0 or threshold <= 0:
+            raise ValueError("invalid Forward-Forward hyperparameters")
+        self.ff_threshold = float(threshold)
+        self.ff_score_from_layer = int(score_from_layer)
+        self.ff_num_layers = self.n_hidden + 1
+        if not 0 <= self.ff_score_from_layer < self.ff_num_layers:
+            raise ValueError("Forward-Forward score range excludes all layers")
+        self.ff_optimizers = []
+        for index, weight in enumerate(self.W):
+            parameters = [weight]
+            if self.normalization == "batchnorm":
+                parameters.extend([self.gamma[index], self.beta[index]])
+            for parameter in parameters:
+                parameter.requires_grad_(True)
+            self.ff_optimizers.append(
+                torch.optim.Adam(parameters, lr=learning_rate))
+        self.W_out.requires_grad_(True)
+        self.b_out.requires_grad_(True)
+        self.ff_optimizers.append(torch.optim.Adam(
+            [self.W_out, self.b_out], lr=learning_rate))
+
+    @staticmethod
+    def _ff_normalize(value):
+        axes = tuple(range(1, value.ndim))
+        norm = torch.sqrt(torch.sum(value.square(), dim=axes, keepdim=True))
+        return value / (norm + 1e-8)
+
+    def ff_overlay(self, image, labels):
+        flat = image.reshape(image.shape[0], -1).clone()
+        flat[:, :self.n_classes] = 0.0
+        flat[torch.arange(image.shape[0], device=image.device), labels] = (
+            torch.max(image).detach())
+        return flat.reshape_as(image)
+
+    def ff_forward(self, image, training=False, update_stats=False):
+        """Forward candidate-labelled data through normalized residual edges."""
+        hiddens = []
+        caches = []
+        pre = self._ff_normalize(image)
+        convolution = F.conv2d(pre, self.W[0], stride=1, padding=1)
+        normalized, norm_cache = self._normalize(
+            0, convolution, training, update_stats)
+        hidden = F.relu(normalized)
+        hiddens.append(hidden)
+        caches.append({
+            "pre": pre, "normalization": norm_cache, "shortcut": None})
+        for block in self.blocks:
+            first = block["first"]
+            second = block["second"]
+            parent = hidden
+            pre = self._ff_normalize(parent)
+            convolution = F.conv2d(
+                pre, self.W[first], stride=block["stride"], padding=1)
+            normalized, norm_cache = self._normalize(
+                first, convolution, training, update_stats)
+            first_hidden = F.relu(normalized)
+            hiddens.append(first_hidden)
+            caches.append({
+                "pre": pre, "normalization": norm_cache, "shortcut": None})
+
+            pre = self._ff_normalize(first_hidden)
+            shortcut = self._option_a_shortcut(
+                parent, block["out_channels"], block["stride"])
+            convolution = F.conv2d(
+                pre, self.W[second], stride=1, padding=1)
+            normalized, norm_cache = self._normalize(
+                second, convolution, training, update_stats)
+            hidden = F.relu(
+                shortcut + self.residual_scale * normalized)
+            hiddens.append(hidden)
+            caches.append({
+                "pre": pre, "normalization": norm_cache,
+                "shortcut": shortcut})
+        features = self._ff_normalize(hidden.mean(dim=(2, 3)))
+        logits = features @ self.W_out.t() + self.b_out
+        return {
+            "hiddens": hiddens,
+            "caches": caches,
+            "features": features,
+            "logits": logits,
+        }
+
+    def _ff_local_outputs(self, layer_index, positive, negative):
+        """Re-evaluate exactly one target layer on detached prefix states."""
+        if layer_index == self.n_hidden:
+            inputs = torch.cat(
+                [positive["features"], negative["features"]], dim=0).detach()
+            outputs = inputs @ self.W_out.t() + self.b_out
+        else:
+            positive_cache = positive["caches"][layer_index]
+            negative_cache = negative["caches"][layer_index]
+            inputs = torch.cat(
+                [positive_cache["pre"], negative_cache["pre"]],
+                dim=0).detach()
+            spec = self.layer_specs[layer_index]
+            convolution = F.conv2d(
+                inputs, self.W[layer_index],
+                stride=spec.stride, padding=spec.padding)
+            normalized, _ = self._normalize(
+                layer_index, convolution, training=True, update_stats=True)
+            if positive_cache["shortcut"] is None:
+                outputs = F.relu(normalized)
+            else:
+                shortcut = torch.cat([
+                    positive_cache["shortcut"],
+                    negative_cache["shortcut"],
+                ], dim=0).detach()
+                outputs = F.relu(
+                    shortcut + spec.branch_scale * normalized)
+        return outputs.chunk(2, dim=0)
+
+    def ff_train_layer(self, layer_index, image, labels,
+                       learning_rate=0.03, negative_labels=None):
+        """Update one greedy FF layer; every prefix and non-target is detached."""
+        if not 0 <= layer_index < self.ff_num_layers:
+            raise ValueError("invalid Forward-Forward layer")
+        if negative_labels is None:
+            offsets = torch.randint(
+                1, self.n_classes, labels.shape, device=labels.device)
+            negative_labels = (labels + offsets) % self.n_classes
+        with torch.no_grad():
+            positive = self.ff_forward(
+                self.ff_overlay(image, labels),
+                training=True, update_stats=False)
+            negative = self.ff_forward(
+                self.ff_overlay(image, negative_labels),
+                training=True, update_stats=False)
+        all_parameters = (
+            self.W + self.gamma + self.beta + [self.W_out, self.b_out])
+        for parameter in all_parameters:
+            parameter.grad = None
+        optimizer = self.ff_optimizers[layer_index]
+        optimizer.param_groups[0]["lr"] = learning_rate
+        positive_output, negative_output = self._ff_local_outputs(
+            layer_index, positive, negative)
+        axes = tuple(range(1, positive_output.ndim))
+        positive_goodness = positive_output.square().mean(dim=axes)
+        negative_goodness = negative_output.square().mean(dim=axes)
+        loss = (
+            F.softplus(-positive_goodness + self.ff_threshold)
+            + F.softplus(negative_goodness - self.ff_threshold)
+        ).mean()
+        loss.backward()
+        optimizer.step()
+        return {
+            "loss": float(loss.detach()),
+            "positive_goodness": float(positive_goodness.mean().detach()),
+            "negative_goodness": float(negative_goodness.mean().detach()),
+            "pair_accuracy": float(
+                (positive_goodness > negative_goodness).float().mean()),
+        }
+
+    @torch.no_grad()
+    def ff_candidate_scores(self, image):
+        scores = []
+        for candidate in range(self.n_classes):
+            labels = torch.full(
+                (image.shape[0],), candidate, device=image.device,
+                dtype=torch.long)
+            forward = self.ff_forward(
+                self.ff_overlay(image, labels), training=False)
+            layer_goodness = [
+                value.square().mean(dim=tuple(range(1, value.ndim)))
+                for value in forward["hiddens"] + [forward["logits"]]]
+            scores.append(sum(
+                layer_goodness[self.ff_score_from_layer:]))
+        return torch.stack(scores, dim=1)
+
+
 class CIFARDualPropResNet(CIFARLocalResNet):
     """Dual Propagation on the residual DAG with author DP-transpose updates.