From e1739ab0a9f1381baff6dd4b9054f8bb7773c868 Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Mon, 27 Jul 2026 14:01:15 -0500 Subject: baseline: add matched ResNet Forward-Forward --- experiments/resnet_crossover_smoke.py | 85 ++++++++++++++++- sdil/conv_crossover.py | 175 ++++++++++++++++++++++++++++++++++ 2 files changed, 259 insertions(+), 1 deletion(-) 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. -- cgit v1.2.3