experiment: package portable A6000 crossover

1 files changed, 422 insertions, 0 deletions

diff --git a/external/dualprop_patches/0005-crossover-add-greedy-Forward-Forward-runner.patch b/external/dualprop_patches/0005-crossover-add-greedy-Forward-Forward-runner.patch
new file mode 100644
index 0000000..1b5f46f
--- /dev/null
+++ b/external/dualprop_patches/0005-crossover-add-greedy-Forward-Forward-runner.patch
@@ -0,0 +1,422 @@
+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
+