path: root/external/dualprop_patches/0005-crossover-add-greedy-Forward-Forward-runner.patch

From fe79d6281891ac52977416bb9d65faac3c32e8e4 Mon Sep 17 00:00:00 2001
From: YurenHao0426 <Blackhao0426@gmail.com>
Date: Mon, 27 Jul 2026 13:12:55 -0500
Subject: [PATCH 05/19] crossover: add greedy Forward-Forward runner

---
 config/cli_config.py       |  12 ++++-
 src/__init__.py            |   2 +-
 src/models.py              |  43 +++++++++++++++
 src/training_utils.py      | 108 +++++++++++++++++++++++++++++++++++++
 tests/local_rules_smoke.py |  40 +++++++++++++-
 train.py                   |   4 ++
 train_ff.py                |  92 +++++++++++++++++++++++++++++++
 7 files changed, 297 insertions(+), 4 deletions(-)
 create mode 100644 train_ff.py

diff --git a/config/cli_config.py b/config/cli_config.py
index 3088436..159e9bc 100644
--- a/config/cli_config.py
+++ b/config/cli_config.py
@@ -42,7 +42,7 @@ parser.add_argument('--experiment-name', default='test', help='A string denoting
 
 parser.add_argument('--model', default='VGG16', choices=['VGG16', 'VGGlike', 'CNN', 'miniCNN', 'MLP'], help='')
 
-parser.add_argument('--learning-algorithm', default='dualprop-lagr-ff', choices=['backprop', 'fa', 'dfa', 'pepita', 'dualprop-lagr-ff', 'dualprop-raovr-ff', 'dualprop-raovr-dampened-ff'])
+parser.add_argument('--learning-algorithm', default='dualprop-lagr-ff', choices=['backprop', 'fa', 'dfa', 'pepita', 'ff', 'dualprop-lagr-ff', 'dualprop-raovr-ff', 'dualprop-raovr-dampened-ff'])
 
 parser.add_argument(
     '--feedback-seed', default=1729, type=int,
@@ -52,6 +52,14 @@ parser.add_argument(
     '--pepita-projection-scale', default=0.05, type=float,
     help='Multiplier on PEPITA He-uniform output-error-to-input projection.')
 
+parser.add_argument(
+    '--ff-threshold', default=2.0, type=float,
+    help='Forward-Forward positive/negative goodness threshold.')
+
+parser.add_argument(
+    '--ff-score-from-layer', default=1, type=int,
+    help='First zero-indexed FF layer included in candidate-label goodness.')
+
 parser.add_argument(
     '--gradient-diagnostics', default='full', choices=['none', 'full'],
     help=('Compute the exact BP reference gradient and layerwise cosine on '
@@ -132,7 +140,7 @@ elif config.model == "MLP":
     dense_features = [1024, 1024, config.num_classes]
 
 # Load model
-modeltype = {"backprop":cnn_abstract, "fa": cnn_abstract, "dfa": cnn_abstract, "pepita": cnn_abstract, "dualprop-lagr-ff": cnn_dualprop_Lagr_ff, "dualprop-raovr-ff": cnn_dualprop_RAOVR_ff, "dualprop-raovr-dampened-ff": cnn_dualprop_RAOVR_dampened_ff}
+modeltype = {"backprop":cnn_abstract, "fa": cnn_abstract, "dfa": cnn_abstract, "pepita": cnn_abstract, "ff": cnn_abstract, "dualprop-lagr-ff": cnn_dualprop_Lagr_ff, "dualprop-raovr-ff": cnn_dualprop_RAOVR_ff, "dualprop-raovr-dampened-ff": cnn_dualprop_RAOVR_dampened_ff}
 activation={"relu": relu, "hs": hs, "sigmoid": sigmoid, "tanh": tanh}
 config.model = modeltype[config.learning_algorithm](loss_func, Conv, Dense, activation[config.activation], config.num_classes, config.beta, config.alpha, config.dtype, config.param_dtype,
                                     kernels=kernels, strides=strides, features=features, mp = mp,
diff --git a/src/__init__.py b/src/__init__.py
index 7a08c87..752dcc2 100644
--- a/src/__init__.py
+++ b/src/__init__.py
@@ -1,2 +1,2 @@
 from .models import cnn_dualprop_Lagr_ff, cnn_dualprop_RAOVR_ff, cnn_dualprop_RAOVR_dampened_ff, cnn_abstract
-from .training_utils import create_train_state, create_local_feedback, train_epoch, eval_model, get_mnist, get_svhn, get_fashionmnist, get_cifar10, get_cifar100, get_imagenet_32x32, heatmap_grads_batches, heatmap_grads_epochs, plot_L_or_gamma
+from .training_utils import create_train_state, create_ff_train_state, create_local_feedback, train_epoch, train_ff_epoch, eval_model, eval_ff_model, get_mnist, get_svhn, get_fashionmnist, get_cifar10, get_cifar100, get_imagenet_32x32, heatmap_grads_batches, heatmap_grads_epochs, plot_L_or_gamma
diff --git a/src/models.py b/src/models.py
index 127268a..0f5af4f 100644
--- a/src/models.py
+++ b/src/models.py
@@ -231,6 +231,49 @@ class cnn_abstract(nn.Module, ABC):
                 prediction = layer(x)
                 objective += jnp.sum(prediction * field) / batch_size
         return objective
+
+    @staticmethod
+    def _ff_normalize(x):
+        axes = tuple(range(1, x.ndim))
+        norm = jnp.sqrt(jnp.sum(jnp.square(x), axis=axes, keepdims=True))
+        return x / (norm + 1e-8)
+
+    def _ff_layer(self, x, layer_index):
+        x = self._ff_normalize(x)
+        return self.act(self.layers[layer_index](x))
+
+    def ff_prefix(self, x, stop_layer):
+        """Detached input producer for greedy Forward-Forward training."""
+        for i in range(stop_layer):
+            x = self._ff_layer(x, i)
+        return x
+
+    def ff_layer_objective(self, positive_input, negative_input, layer_index,
+                           threshold):
+        """Reference Forward-Forward local softplus goodness objective."""
+        positive = self._ff_layer(
+            jax.lax.stop_gradient(positive_input), layer_index)
+        negative = self._ff_layer(
+            jax.lax.stop_gradient(negative_input), layer_index)
+        axes = tuple(range(1, positive.ndim))
+        positive_goodness = jnp.mean(jnp.square(positive), axis=axes)
+        negative_goodness = jnp.mean(jnp.square(negative), axis=axes)
+        loss = jnp.mean(
+            jax.nn.softplus(-positive_goodness + threshold)
+            + jax.nn.softplus(negative_goodness - threshold))
+        pair_accuracy = jnp.mean(
+            positive_goodness > negative_goodness, dtype=jnp.float32)
+        return loss, (positive_goodness, negative_goodness, pair_accuracy)
+
+    def ff_goodness(self, x, score_from_layer=1):
+        """Candidate-label goodness used by supervised FF inference."""
+        score = jnp.zeros((x.shape[0],), dtype=x.dtype)
+        for i in range(self.num_layers):
+            x = self._ff_layer(x, i)
+            if i >= score_from_layer:
+                axes = tuple(range(1, x.ndim))
+                score += jnp.mean(jnp.square(x), axis=axes)
+        return score
     
     def init_states_to_zero(self, x0):
         s = [x0]
diff --git a/src/training_utils.py b/src/training_utils.py
index f5fae62..734f46d 100644
--- a/src/training_utils.py
+++ b/src/training_utils.py
@@ -218,6 +218,15 @@ def create_train_state(rng, model, image_dims, lr, wlr, lrf, momentum, weight_de
     return train_state.TrainState.create(apply_fn=model.apply, params=unfreeze(params), tx=tx)
 
 
+def create_ff_train_state(rng, model, image_dims, learning_rate):
+    """Reference-style Adam state for greedy Forward-Forward layers."""
+    w, h, channels = image_dims
+    dummy = jnp.ones([1, w, h, channels])
+    params = unfreeze(model.init(rng, dummy)["params"])
+    return train_state.TrainState.create(
+        apply_fn=model.apply, params=params, tx=optax.adam(learning_rate))
+
+
 def create_local_feedback(rng, model, params, image_dims, learning_algorithm,
                           num_classes, pepita_projection_scale=0.05):
     """Create fixed feedback without reading a forward parameter value.
@@ -342,6 +351,105 @@ def direct_feedback_fields(states, output_field, direct_feedback):
     return fields
 
 
+def ff_overlay(image, labels, num_classes):
+    """Overlay a candidate label on the first input coordinates."""
+    flat = image.reshape((image.shape[0], -1))
+    flat = flat.at[:, :num_classes].set(0.0)
+    flat = flat.at[jnp.arange(image.shape[0]), labels].set(jnp.max(image))
+    return flat.reshape(image.shape)
+
+
+@partial(jax.jit, static_argnames=("layer_index", "num_classes"))
+def train_step_ff(state, image, labels, batch_rng, augmentation_on,
+                  layer_index, num_classes, threshold):
+    """One greedy Forward-Forward update of exactly one layer."""
+    augmentation_rng, negative_rng = jax.random.split(batch_rng)
+    per_example_rng = jax.random.split(augmentation_rng, image.shape[0])
+    image = jax.lax.cond(
+        augmentation_on, vmap_augment_train, no_aug, image, per_example_rng)
+    offsets = jax.random.randint(
+        negative_rng, labels.shape, 1, num_classes)
+    negative_labels = (labels + offsets) % num_classes
+    positive = ff_overlay(image, labels, num_classes)
+    negative = ff_overlay(image, negative_labels, num_classes)
+    positive_input = state.apply_fn(
+        {"params": state.params}, positive, layer_index, method="ff_prefix")
+    negative_input = state.apply_fn(
+        {"params": state.params}, negative, layer_index, method="ff_prefix")
+    positive_input = jax.lax.stop_gradient(positive_input)
+    negative_input = jax.lax.stop_gradient(negative_input)
+
+    def objective(params):
+        return state.apply_fn(
+            {"params": params}, positive_input, negative_input, layer_index,
+            threshold, method="ff_layer_objective")
+
+    (loss, auxiliary), grads = jax.value_and_grad(
+        objective, has_aux=True)(state.params)
+    state = state.apply_gradients(grads=grads)
+    return state, {
+        "loss": loss,
+        "positive_goodness": jnp.mean(auxiliary[0]),
+        "negative_goodness": jnp.mean(auxiliary[1]),
+        "pair_accuracy": auxiliary[2],
+    }
+
+
+def train_ff_epoch(state, train_ds, batch_size, rng, augmentation_on,
+                   layer_index, num_classes, threshold):
+    """Train one FF layer for one full data epoch."""
+    t0 = time.time()
+    train_ds_size = len(train_ds["image"])
+    steps_per_epoch = train_ds_size // batch_size
+    perms = jax.random.permutation(rng, train_ds_size)
+    perms = perms[:steps_per_epoch * batch_size]
+    perms = perms.reshape((steps_per_epoch, batch_size))
+    metrics = []
+    for permutation in perms:
+        image = train_ds["image"][permutation]
+        labels = train_ds["label"][permutation]
+        rng, batch_rng = jax.random.split(rng)
+        state, batch_metrics = train_step_ff(
+            state, image, labels, batch_rng, augmentation_on, layer_index,
+            num_classes, threshold)
+        metrics.append(batch_metrics)
+    host_metrics = jax.device_get(metrics)
+    summary = {
+        key: float(np.mean([record[key] for record in host_metrics]))
+        for key in host_metrics[0]
+    }
+    return state, summary, time.time() - t0
+
+
+@partial(jax.jit, static_argnames=("num_classes", "score_from_layer"))
+def predict_ff(state, image, num_classes, score_from_layer):
+    scores = []
+    for candidate in range(num_classes):
+        labels = jnp.full((image.shape[0],), candidate, dtype=jnp.int32)
+        overlaid = ff_overlay(image, labels, num_classes)
+        scores.append(state.apply_fn(
+            {"params": state.params}, overlaid, score_from_layer,
+            method="ff_goodness"))
+    return jnp.stack(scores, axis=-1)
+
+
+def eval_ff_model(state, dataset, batch_size, num_classes, score_from_layer):
+    """Evaluate all candidate-label overlays; no classifier head is assumed."""
+    t0 = time.time()
+    size = len(dataset["image"])
+    steps = size // batch_size
+    indices = jnp.arange(steps * batch_size).reshape((steps, batch_size))
+    correct = 0
+    total = 0
+    for index in indices:
+        image = dataset["image"][index]
+        labels = dataset["label"][index]
+        scores = predict_ff(state, image, num_classes, score_from_layer)
+        correct += int(jax.device_get(jnp.sum(jnp.argmax(scores, -1) == labels)))
+        total += labels.shape[0]
+    return 100.0 * correct / total, time.time() - t0
+
+
 @jax.jit
 def train_step_pepita(state, input_feedback, image, labels_onehot, labels,
                       batch_rng, augmentation_on, gradient_diagnostics):
diff --git a/tests/local_rules_smoke.py b/tests/local_rules_smoke.py
index e7da52e..8ce1e02 100644
--- a/tests/local_rules_smoke.py
+++ b/tests/local_rules_smoke.py
@@ -15,7 +15,11 @@ ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 sys.path.insert(0, ROOT)
 
 from src.models import cnn_abstract
-from src.training_utils import create_local_feedback, direct_feedback_fields
+from src.training_utils import (
+    create_local_feedback,
+    direct_feedback_fields,
+    ff_overlay,
+)
 
 
 def loss_func(logits, one_hot):
@@ -153,6 +157,38 @@ def main():
     assert float(jnp.max(jnp.abs(pepita_projection))) <= projection_limit
     assert bool(jnp.all(jnp.isfinite(flat(pepita_grads))))
 
+    positive = ff_overlay(x, labels, 3)
+    negative = ff_overlay(x, (labels + 1) % 3, 3)
+    assert bool(jnp.all(positive.reshape((x.shape[0], -1))[
+        jnp.arange(x.shape[0]), labels] == jnp.max(x)))
+    positive_prefix = model.apply(
+        {"params": params}, positive, 1, method="ff_prefix")
+    negative_prefix = model.apply(
+        {"params": params}, negative, 1, method="ff_prefix")
+    (ff_loss, ff_aux), ff_grads = jax.value_and_grad(
+        lambda candidate: model.apply(
+            {"params": candidate}, positive_prefix, negative_prefix, 1, 2.0,
+            method="ff_layer_objective"),
+        has_aux=True,
+    )(params)
+    ff_manual_loss = jnp.mean(
+        jax.nn.softplus(-ff_aux[0] + 2.0)
+        + jax.nn.softplus(ff_aux[1] - 2.0))
+    ff_objective_error = float(jnp.abs(ff_loss - ff_manual_loss))
+    assert ff_objective_error < 2e-12
+    assert float(jnp.linalg.norm(flat(ff_grads["c01"]))) > 0.0
+    assert float(jnp.linalg.norm(flat(ff_grads["c00"]))) == 0.0
+    assert float(jnp.linalg.norm(flat(ff_grads["d00"]))) == 0.0
+    ff_scores = jnp.stack([
+        model.apply(
+            {"params": params},
+            ff_overlay(x, jnp.full(labels.shape, candidate), 3), 1,
+            method="ff_goodness")
+        for candidate in range(3)
+    ], axis=-1)
+    assert ff_scores.shape == (x.shape[0], 3)
+    assert bool(jnp.all(jnp.isfinite(ff_scores)))
+
     report = {
         "symmetric_fa_bp_relative_error": symmetric_error,
         "independent_feedback_forward_cosine": feedback_cosine,
@@ -161,6 +197,8 @@ def main():
         "dfa_hidden_maps": len(direct),
         "pepita_readout_equation_max_error": pepita_readout_error,
         "pepita_projection_shape": list(pepita_projection.shape),
+        "ff_local_objective_error": ff_objective_error,
+        "ff_score_shape": list(ff_scores.shape),
         "status": "passed",
     }
     print(json.dumps(report, indent=2, sort_keys=True))
diff --git a/train.py b/train.py
index ba5f6ac..d3f0de6 100644
--- a/train.py
+++ b/train.py
@@ -14,6 +14,10 @@ from src import create_train_state, create_local_feedback, train_epoch, eval_mod
 # import config # Use this for the old method
 from config.cli_config import config
 
+if config.learning_algorithm == "ff":
+    raise ValueError(
+        "Forward-Forward uses greedy layerwise training; run train_ff.py")
+
 experiment_dir = "./runs/" + config.experiment_name + "/"
 if experiment_dir == "./runs/debug-test/" and os.path.isdir(experiment_dir):
     shutil.rmtree(experiment_dir)
diff --git a/train_ff.py b/train_ff.py
new file mode 100644
index 0000000..3839dd3
--- /dev/null
+++ b/train_ff.py
@@ -0,0 +1,92 @@
+"""Greedy supervised Forward-Forward on the author plain-CNN topologies."""
+import datetime
+import os
+import time
+
+import jax
+import numpy as np
+
+from config.cli_config import config
+from src import create_ff_train_state, eval_ff_model, train_ff_epoch
+
+
+if config.learning_algorithm != "ff":
+    raise ValueError("train_ff.py requires --learning-algorithm ff")
+if config.ff_score_from_layer < 0:
+    raise ValueError("--ff-score-from-layer must be nonnegative")
+
+experiment_dir = os.path.join("runs", config.experiment_name)
+if os.path.isdir(experiment_dir):
+    raise FileExistsError(
+        "experiment directory exists; refusing to overwrite "
+        + experiment_dir)
+os.makedirs(experiment_dir)
+
+for experiment_index, seed in enumerate(config.seeds):
+    timestamp = datetime.datetime.fromtimestamp(time.time())
+    outpath = os.path.join(
+        experiment_dir, timestamp.strftime("%Y_%m_%d_%H_%M_%S"))
+    os.makedirs(outpath)
+    print(
+        f"Starting FF experiment {experiment_index + 1}/"
+        f"{len(config.seeds)} seed={seed}",
+        flush=True)
+    rng = jax.random.PRNGKey(seed)
+    rng, init_rng = jax.random.split(rng)
+    state = create_ff_train_state(
+        init_rng, config.model, config.image_dims, config.learning_rate)
+    num_layers = len(state.params)
+    if config.ff_score_from_layer >= num_layers:
+        raise ValueError("--ff-score-from-layer excludes every layer")
+    augmentation_on = config.dataset not in ("mnist", "fashionmnist")
+    history = {
+        "method": "ff",
+        "epochs_per_layer": config.num_epochs,
+        "num_layers": num_layers,
+        "threshold": config.ff_threshold,
+        "score_from_layer": config.ff_score_from_layer,
+        "learning_rate": config.learning_rate,
+        "layers": [],
+    }
+    started = time.time()
+    for layer_index in range(num_layers):
+        layer_record = {"layer": layer_index, "epochs": []}
+        for epoch in range(config.num_epochs):
+            rng, epoch_rng = jax.random.split(rng)
+            state, metrics, runtime = train_ff_epoch(
+                state, config.train_ds, config.batch_size, epoch_rng,
+                augmentation_on, layer_index, config.num_classes,
+                config.ff_threshold)
+            record = {
+                "epoch": epoch + 1,
+                "runtime": runtime,
+                **metrics,
+            }
+            layer_record["epochs"].append(record)
+            print(
+                f"layer={layer_index + 1}/{num_layers} "
+                f"epoch={epoch + 1}/{config.num_epochs} "
+                f"loss={metrics['loss']:.5f} "
+                f"pair_acc={100 * metrics['pair_accuracy']:.2f}% "
+                f"runtime={runtime:.3f}s",
+                flush=True)
+        history["layers"].append(layer_record)
+
+    validation_accuracy, validation_time = eval_ff_model(
+        state, config.val_ds, config.batch_size, config.num_classes,
+        config.ff_score_from_layer)
+    test_accuracy, test_time = eval_ff_model(
+        state, config.test_ds, config.batch_size, config.num_classes,
+        config.ff_score_from_layer)
+    history["final"] = {
+        "validation_accuracy": validation_accuracy,
+        "validation_time": validation_time,
+        "test_accuracy": test_accuracy,
+        "test_time": test_time,
+        "train_and_eval_wall": time.time() - started,
+    }
+    print(
+        f"final val_accuracy={validation_accuracy:.3f}% "
+        f"test_accuracy={test_accuracy:.3f}%",
+        flush=True)
+    np.save(os.path.join(outpath, "hist.npy"), history)
-- 
2.54.0