From fe79d6281891ac52977416bb9d65faac3c32e8e4 Mon Sep 17 00:00:00 2001 From: YurenHao0426 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