diff options
Diffstat (limited to 'external/dualprop_patches/0004-crossover-add-architecture-compatible-PEPITA.patch')
| -rw-r--r-- | external/dualprop_patches/0004-crossover-add-architecture-compatible-PEPITA.patch | 242 |
1 files changed, 242 insertions, 0 deletions
diff --git a/external/dualprop_patches/0004-crossover-add-architecture-compatible-PEPITA.patch b/external/dualprop_patches/0004-crossover-add-architecture-compatible-PEPITA.patch new file mode 100644 index 0000000..896350b --- /dev/null +++ b/external/dualprop_patches/0004-crossover-add-architecture-compatible-PEPITA.patch @@ -0,0 +1,242 @@ +From 3bfbe71fcdc138764408a5e1646b66c464308b8b Mon Sep 17 00:00:00 2001 +From: YurenHao0426 <Blackhao0426@gmail.com> +Date: Mon, 27 Jul 2026 13:09:05 -0500 +Subject: [PATCH 04/19] crossover: add architecture-compatible PEPITA + +--- + config/cli_config.py | 8 ++++-- + src/models.py | 37 ++++++++++++++++++++++++++++ + src/training_utils.py | 50 ++++++++++++++++++++++++++++++++++++-- + tests/local_rules_smoke.py | 32 ++++++++++++++++++++++++ + train.py | 3 ++- + 5 files changed, 125 insertions(+), 5 deletions(-) + +diff --git a/config/cli_config.py b/config/cli_config.py +index e1756e1..3088436 100644 +--- a/config/cli_config.py ++++ b/config/cli_config.py +@@ -42,12 +42,16 @@ 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', 'dualprop-lagr-ff', 'dualprop-raovr-ff', 'dualprop-raovr-dampened-ff']) ++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( + '--feedback-seed', default=1729, type=int, + help='Independent fixed-feedback initialization seed for FA/DFA.') + ++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( + '--gradient-diagnostics', default='full', choices=['none', 'full'], + help=('Compute the exact BP reference gradient and layerwise cosine on ' +@@ -128,7 +132,7 @@ elif config.model == "MLP": + dense_features = [1024, 1024, config.num_classes] + + # Load model +-modeltype = {"backprop":cnn_abstract, "fa": cnn_abstract, "dfa": 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, "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/models.py b/src/models.py +index edc432c..127268a 100644 +--- a/src/models.py ++++ b/src/models.py +@@ -194,6 +194,43 @@ class cnn_abstract(nn.Module, ABC): + _, linear_pullback = jax.vjp(linear_map, s[i]) + fields[i] = linear_pullback(linear_field)[0] + return fields ++ ++ def pepita_correlation_objective(self, modulated_states, original_linear, ++ modulated_linear, one_hot): ++ """Architecture-compatible PEPITA/ERIN two-presentation update. ++ ++ Hidden fields are the first-minus-second post-ReLU activity ++ differences from the PEPITA equation. They multiply the modulated ++ presynaptic state directly, without differentiating a nonlinearity. ++ Convolutional correlations follow the reference code's additional ++ average over spatial positions. The readout uses the second-pass ++ softmax error. ++ """ ++ objective = 0.0 ++ batch_size = modulated_states[0].shape[0] ++ for i, layer in enumerate(self.layers): ++ x = jax.lax.stop_gradient(modulated_states[i]) ++ if i == self.num_layers - 1: ++ field = jax.lax.stop_gradient( ++ jax.nn.softmax(modulated_states[-1], axis=-1) - one_hot) ++ prediction = layer(x) ++ objective += jnp.sum(prediction * field) / batch_size ++ elif i < self.num_convlayers: ++ field = jax.lax.stop_gradient( ++ self.act(original_linear[i]) ++ - self.act(modulated_linear[i])) ++ prediction = layer.call_without_pooling(x) ++ spatial_positions = prediction.shape[1] * prediction.shape[2] ++ objective += ( ++ jnp.sum(prediction * field) ++ / (batch_size * spatial_positions)) ++ else: ++ field = jax.lax.stop_gradient( ++ self.act(original_linear[i]) ++ - self.act(modulated_linear[i])) ++ prediction = layer(x) ++ objective += jnp.sum(prediction * field) / batch_size ++ return objective + + def init_states_to_zero(self, x0): + s = [x0] +diff --git a/src/training_utils.py b/src/training_utils.py +index 2b6b18a..f5fae62 100644 +--- a/src/training_utils.py ++++ b/src/training_utils.py +@@ -219,17 +219,23 @@ def create_train_state(rng, model, image_dims, lr, wlr, lrf, momentum, weight_de + + + def create_local_feedback(rng, model, params, image_dims, learning_algorithm, +- num_classes): ++ num_classes, pepita_projection_scale=0.05): + """Create fixed feedback without reading a forward parameter value. + + FA uses an independently initialized model-shaped linear parameter tree. + DFA uses one independent output-to-hidden tensor per hidden population. + The forward tree is used only to infer activation shapes for DFA. + """ +- if learning_algorithm not in ("fa", "dfa"): ++ if learning_algorithm not in ("fa", "dfa", "pepita"): + return None + w, h, channels = image_dims + dummy = jnp.ones([1, w, h, channels]) ++ if learning_algorithm == "pepita": ++ input_units = w * h * channels ++ limit = jnp.sqrt(6.0 / input_units) * pepita_projection_scale ++ return jax.random.uniform( ++ rng, (num_classes, w, h, channels), ++ minval=-limit, maxval=limit, dtype=dummy.dtype) + if learning_algorithm == "fa": + return unfreeze(model.init(rng, dummy)["params"]) + +@@ -301,6 +307,10 @@ def train_epoch(state, train_ds, batch_size, rng, augmentation_on, + state, local_feedback, learning_algorithm, image, + labels_onehot, labels, batch_rng, augmentation_on, + gradient_diagnostics) ++ elif learning_algorithm == "pepita": ++ state, metrics = train_step_pepita( ++ state, local_feedback, image, labels_onehot, labels, ++ batch_rng, augmentation_on, gradient_diagnostics) + elif learning_algorithm != "backprop": + state, metrics = train_step( + state, image, labels_onehot, labels, batch_rng, inf_rng, +@@ -332,6 +342,42 @@ def direct_feedback_fields(states, output_field, direct_feedback): + return fields + + ++@jax.jit ++def train_step_pepita(state, input_feedback, image, labels_onehot, labels, ++ batch_rng, augmentation_on, gradient_diagnostics): ++ """One PEPITA step: ordinary forward, error-modulated forward, local rule.""" ++ image = jax.lax.cond( ++ augmentation_on, vmap_augment_train, no_aug, image, batch_rng) ++ states, linear = state.apply_fn( ++ {"params": state.params}, image, method="ff_with_local_cache") ++ output_error = jax.lax.stop_gradient( ++ jax.nn.softmax(states[-1], axis=-1) - labels_onehot) ++ input_error = jnp.tensordot( ++ output_error, input_feedback, axes=((-1,), (0,))) ++ modulated_image = jax.lax.stop_gradient(image + input_error) ++ modulated_states, modulated_linear = state.apply_fn( ++ {"params": state.params}, modulated_image, ++ method="ff_with_local_cache") ++ states = tree_map(jax.lax.stop_gradient, states) ++ linear = tree_map(jax.lax.stop_gradient, linear) ++ modulated_states = tree_map(jax.lax.stop_gradient, modulated_states) ++ modulated_linear = tree_map(jax.lax.stop_gradient, modulated_linear) ++ ++ def local_objective(params): ++ return state.apply_fn( ++ {"params": params}, modulated_states, linear, modulated_linear, ++ labels_onehot, method="pepita_correlation_objective") ++ ++ grads = jax.grad(local_objective)(state.params) ++ metrics = compute_metrics( ++ image=image, labels_onehot=labels_onehot, labels=labels, state=state) ++ metrics = jax.lax.cond( ++ gradient_diagnostics, ref_grad_and_angle, no_ref_grad_and_angle, ++ state, grads, image, labels_onehot, metrics) ++ state = state.apply_gradients(grads=grads) ++ return state, metrics ++ ++ + @partial(jax.jit, static_argnames=("learning_algorithm",)) + def train_step_local(state, local_feedback, learning_algorithm, image, + labels_onehot, labels, batch_rng, augmentation_on, +diff --git a/tests/local_rules_smoke.py b/tests/local_rules_smoke.py +index 2366e61..e7da52e 100644 +--- a/tests/local_rules_smoke.py ++++ b/tests/local_rules_smoke.py +@@ -123,12 +123,44 @@ def main(): + ) + assert all(bool(jnp.all(jnp.isfinite(field))) for field in dfa_fields) + ++ pepita_projection = create_local_feedback( ++ jax.random.PRNGKey(1731), model, params, (8, 8, 2), "pepita", 3, ++ pepita_projection_scale=0.05) ++ pepita_error = jax.nn.softmax(states[-1], axis=-1) - one_hot ++ modulated_x = x + jnp.tensordot( ++ pepita_error, pepita_projection, axes=((-1,), (0,))) ++ modulated_states, modulated_linear = model.apply( ++ {"params": params}, modulated_x, method="ff_with_local_cache") ++ pepita_grads = jax.grad( ++ lambda candidate: model.apply( ++ {"params": candidate}, modulated_states, linear, ++ modulated_linear, one_hot, ++ method="pepita_correlation_objective") ++ )(params) ++ output_error = jax.nn.softmax(modulated_states[-1], axis=-1) - one_hot ++ output_input = modulated_states[-2].reshape((x.shape[0], -1)) ++ manual_kernel = output_input.T @ output_error / x.shape[0] ++ manual_bias = output_error.mean(axis=0) ++ readout_leaves = jax.tree_util.tree_leaves(pepita_grads["d00"]) ++ readout_kernel = next(value for value in readout_leaves if value.ndim == 2) ++ readout_bias = next(value for value in readout_leaves if value.ndim == 1) ++ pepita_readout_error = max( ++ float(jnp.max(jnp.abs(readout_kernel - manual_kernel))), ++ float(jnp.max(jnp.abs(readout_bias - manual_bias))), ++ ) ++ assert pepita_readout_error < 2e-12, pepita_readout_error ++ projection_limit = (6.0 / (8 * 8 * 2)) ** 0.5 * 0.05 ++ assert float(jnp.max(jnp.abs(pepita_projection))) <= projection_limit ++ assert bool(jnp.all(jnp.isfinite(flat(pepita_grads)))) ++ + report = { + "symmetric_fa_bp_relative_error": symmetric_error, + "independent_feedback_forward_cosine": feedback_cosine, + "feedback_independence_error": feedback_independence_error, + "detached_local_boundary_error": local_boundary_error, + "dfa_hidden_maps": len(direct), ++ "pepita_readout_equation_max_error": pepita_readout_error, ++ "pepita_projection_shape": list(pepita_projection.shape), + "status": "passed", + } + print(json.dumps(report, indent=2, sort_keys=True)) +diff --git a/train.py b/train.py +index c030158..ba5f6ac 100644 +--- a/train.py ++++ b/train.py +@@ -48,7 +48,8 @@ for experiment_index, seed in enumerate(config.seeds): + state = create_train_state(init_rng, config.model, config.image_dims, config.learning_rate, config.warmup_learning_rate, config.learning_rate_final, config.momentum, config.weight_decay, config.num_epochs, config.warmup_epochs, config.decay_epochs, steps_per_epoch) + local_feedback = create_local_feedback( + jax.random.PRNGKey(config.feedback_seed), config.model, state.params, +- config.image_dims, config.learning_algorithm, config.num_classes) ++ config.image_dims, config.learning_algorithm, config.num_classes, ++ pepita_projection_scale=config.pepita_projection_scale) + + + del init_rng # Must not be used anymore. +-- +2.54.0 + |