"""Matched non-backprop adapters for the audited CIFAR ResNet topology. These adapters reuse :class:`CIFARLocalResNet`'s forward tensors, BatchNorm, option-A shortcuts, and local optimizer. They are kept separate from the established SDIL/KP implementation so a crossover baseline cannot silently change the already confirmed forward model. """ import torch import torch.nn.functional as F from .conv import CIFARLocalResNet class CIFARPEPITAResNet(CIFARLocalResNet): """Two-presentation PEPITA on the matched residual forward topology.""" def __init__(self, *args, projection_scale=0.05, projection_seed=1731, **kwargs): super().__init__(*args, **kwargs) if projection_scale <= 0: raise ValueError("PEPITA projection scale must be positive") generator = torch.Generator(device="cpu").manual_seed(projection_seed) input_units = 3 * 32 * 32 limit = (6.0 / input_units) ** 0.5 * projection_scale projection = ( 2.0 * torch.rand( self.n_classes, 3, 32, 32, generator=generator) - 1.0 ) * limit self.input_feedback = projection.to( device=self.device, dtype=self.dtype) @property def n_fixed_feedback_parameters(self): return self.input_feedback.numel() @torch.no_grad() def pepita_ascent_directions(self, clean, modulated, modulated_error): """Return the explicit first-minus-second PEPITA correlations. The post-activation difference directly multiplies the modulated presynaptic activity, as in PEPITA/ERIN; it is not differentiated through the ReLU. BatchNorm's current-layer Jacobian and a residual branch's fixed multiplier remain part of that local synaptic eligibility. Convolutional correlations follow the reference code's additional average over spatial positions. """ batch = modulated_error.shape[0] directions = [] gamma_directions = [] beta_directions = [] for index, (clean_hidden, modulated_hidden, cache, weight, spec) in ( enumerate(zip( clean["hiddens"], modulated["hiddens"], modulated["caches"], self.W, self.layer_specs))): field = clean_hidden - modulated_hidden local_field = field * spec.branch_scale local_field, gamma_direction, beta_direction = ( self._normalization_backward( index, local_field, cache["normalization"])) spatial = local_field.shape[2] * local_field.shape[3] correlation = torch.nn.grad.conv2d_weight( cache["pre"], weight.shape, local_field, stride=spec.stride, padding=spec.padding) directions.append(-correlation / (batch * spatial)) if gamma_direction is not None: gamma_directions.append( -gamma_direction / (batch * spatial)) beta_directions.append( -beta_direction / (batch * spatial)) output_weight = -( modulated_error.t() @ modulated["features"]) / batch output_bias = -modulated_error.mean(dim=0) return ( directions, gamma_directions, beta_directions, output_weight, output_bias) def pepita_step(self, image, labels, eta, eta_output=None, momentum=0.0, weight_decay=0.0): """One architecture-compatible PEPITA/ERIN local update.""" one_hot = F.one_hot(labels, self.n_classes).to(image.dtype) with torch.no_grad(): clean = self.forward( image, return_cache=True, training=True, update_stats=True) clean_error = torch.softmax(clean["logits"], dim=1) - one_hot input_error = torch.einsum( "bc,cijk->bijk", clean_error, self.input_feedback) modulated = self.forward( image + input_error, return_cache=True, training=True, update_stats=False) modulated_error = ( torch.softmax(modulated["logits"], dim=1) - one_hot) (directions, gamma_directions, beta_directions, output_weight, output_bias) = self.pepita_ascent_directions( clean, modulated, modulated_error) self.apply_ascent( directions, output_weight, output_bias, eta, eta_output=eta_output, momentum=momentum, weight_decay=weight_decay, gamma_directions=gamma_directions, beta_directions=beta_directions) loss = F.cross_entropy(clean["logits"], labels) 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. ``s_plus`` and ``s_minus`` are initialized to the ordinary forward states. A ``fwK`` pass updates every residual-DAG node in topological order. The feedforward drive uses the forward edge, while the difference of each child state is transported through the exact transpose of that same edge. This is intentional symmetric feedback in the Dual Propagation baseline, not a claim of weight-transport-free learning. """ def __init__(self, *args, alpha=0.0, dp_beta=0.1, inference_passes=16, **kwargs): super().__init__(*args, **kwargs) if not 0.0 <= alpha <= 1.0: raise ValueError("Dual Propagation alpha must lie in [0, 1]") if dp_beta <= 0 or inference_passes < 1: raise ValueError("invalid Dual Propagation inference settings") self.dp_alpha = float(alpha) self.dp_beta = float(dp_beta) self.dp_inference_passes = int(inference_passes) self._node_kind = {0: ("stem",)} self._outgoing = {index: [] for index in range(self.n_hidden)} for block in self.blocks: first = block["first"] second = block["second"] parent = first - 1 self._node_kind[first] = ("first", parent) self._node_kind[second] = ( "second", first, parent, block["out_channels"], block["stride"]) self._outgoing[parent].append((first, "conv")) self._outgoing[first].append((second, "conv")) self._outgoing[parent].append((second, "shortcut")) @staticmethod def _option_a_shortcut_transpose(value, input_shape, stride): """Adjoint of the parameter-free option-A shortcut.""" in_channels = input_shape[1] out_channels = value.shape[1] missing = out_channels - in_channels if missing < 0: raise ValueError("option-A transpose cannot increase input width") before = missing // 2 selected = value[:, before:before + in_channels] if stride == 1: if tuple(selected.shape) != tuple(input_shape): raise ValueError("option-A transpose shape mismatch") return selected result = value.new_zeros(input_shape) result[:, :, ::2, ::2] = selected return result def _node_prediction(self, index, states, image): """Return the unrectified local prediction and its edge cache.""" kind = self._node_kind[index] if kind[0] == "stem": pre = image shortcut = None elif kind[0] == "first": pre = states[kind[1]] shortcut = None else: pre = states[kind[1]] shortcut = self._option_a_shortcut( states[kind[2]], kind[3], kind[4]) spec = self.layer_specs[index] convolution = F.conv2d( pre, self.W[index], stride=spec.stride, padding=spec.padding) normalized, normalization = self._normalize( index, convolution, training=True, update_stats=False) prediction = ( shortcut + spec.branch_scale * normalized if shortcut is not None else normalized) return prediction, { "pre": pre, "normalization": normalization, "shortcut": shortcut, } def _edge_transpose(self, child, edge_kind, field, states, image): """Apply one residual-DAG edge transpose to a child state field.""" kind = self._node_kind[child] if edge_kind == "shortcut": if kind[0] != "second": raise AssertionError("only a second convolution has a shortcut") return self._option_a_shortcut_transpose( field, states[kind[2]].shape, kind[4]) _, cache = self._node_prediction(child, states, image) spec = self.layer_specs[child] local_field = field * spec.branch_scale local_field, _, _ = self._normalization_backward( child, local_field, cache["normalization"]) return torch.nn.grad.conv2d_input( cache["pre"].shape, self.W[child], local_field, stride=spec.stride, padding=spec.padding) def _outgoing_feedback(self, index, deltas, states, image): result = torch.zeros_like(states[index]) for child, edge_kind in self._outgoing[index]: result.add_(self._edge_transpose( child, edge_kind, deltas[child], states, image)) if index == self.n_hidden - 1: spatial = states[index].shape[2] * states[index].shape[3] result.add_( (deltas[-1] @ self.W_out)[:, :, None, None] / spatial) return result def infer_dual_states(self, image, one_hot, clean_forward=None): """Run the author ``fwK`` DP-transpose state updates on the DAG.""" if clean_forward is None: clean_forward = self.forward( image, return_cache=True, training=True, update_stats=False) plus = [ value.detach().clone() for value in clean_forward["hiddens"]] minus = [value.detach().clone() for value in plus] plus.append(clean_forward["logits"].detach().clone()) minus.append(clean_forward["logits"].detach().clone()) fixed_prediction = clean_forward["logits"].detach() alpha = self.dp_alpha for _ in range(self.dp_inference_passes): for index in range(self.n_hidden): states = [ alpha * positive + (1.0 - alpha) * negative for positive, negative in zip(plus[:-1], minus[:-1])] prediction, _ = self._node_prediction(index, states, image) deltas = [ positive - negative for positive, negative in zip(plus, minus)] feedback = self._outgoing_feedback( index, deltas, states, image) plus[index] = F.relu( prediction + (1.0 - alpha) * feedback) minus[index] = F.relu(prediction - alpha * feedback) states = [ alpha * positive + (1.0 - alpha) * negative for positive, negative in zip(plus[:-1], minus[:-1])] features = states[-1].mean(dim=(2, 3)) prediction = features @ self.W_out.t() + self.b_out output_field = self.dp_beta * ( torch.softmax(fixed_prediction, dim=1) - one_hot) plus[-1] = prediction - (1.0 - alpha) * output_field minus[-1] = prediction + alpha * output_field return plus, minus @torch.no_grad() def dualprop_ascent_directions(self, image, plus, minus): """Evaluate the local DP contrastive correlations without autograd.""" alpha = self.dp_alpha beta = self.dp_beta states = [ alpha * positive + (1.0 - alpha) * negative for positive, negative in zip(plus[:-1], minus[:-1])] deltas = [ (positive - negative) / beta for positive, negative in zip(plus, minus)] directions = [] gamma_directions = [] beta_directions = [] batch = image.shape[0] for index in range(self.n_hidden): _, cache = self._node_prediction(index, states, image) spec = self.layer_specs[index] local_field = deltas[index] * spec.branch_scale local_field, gamma_direction, beta_direction = ( self._normalization_backward( index, local_field, cache["normalization"])) direction = torch.nn.grad.conv2d_weight( cache["pre"], self.W[index].shape, local_field, stride=spec.stride, padding=spec.padding) directions.append(direction / batch) if gamma_direction is not None: gamma_directions.append(gamma_direction / batch) beta_directions.append(beta_direction / batch) features = states[-1].mean(dim=(2, 3)) output_weight = deltas[-1].t() @ features / batch output_bias = deltas[-1].mean(dim=0) return ( directions, gamma_directions, beta_directions, output_weight, output_bias) def dualprop_step(self, image, labels, eta, eta_output=None, momentum=0.0, weight_decay=0.0): """One fully local DP-transpose update on a matched ResNet batch.""" one_hot = F.one_hot(labels, self.n_classes).to(image.dtype) with torch.no_grad(): clean = self.forward( image, return_cache=True, training=True, update_stats=True) loss = F.cross_entropy(clean["logits"], labels) plus, minus = self.infer_dual_states( image, one_hot, clean_forward=clean) (directions, gamma_directions, beta_directions, output_weight, output_bias) = self.dualprop_ascent_directions( image, plus, minus) self.apply_ascent( directions, output_weight, output_bias, eta, eta_output=eta_output, momentum=momentum, weight_decay=weight_decay, gamma_directions=gamma_directions, beta_directions=beta_directions) return float(loss) class CIFAREquilibriumPropResNet(CIFARDualPropResNet): """Canonical two-phase EP dynamics on the matched residual energy graph. The Dual-Propagation parent class is used only for its audited residual-edge operators and exact edge transposes. EP itself has one hard-sigmoid state per forward population, synchronous leaky relaxation, random nudge sign, and free-versus-nudged local correlations. """ def __init__(self, *args, ep_beta=0.5, dt=0.5, free_steps=20, nudge_steps=4, random_beta_sign=True, **kwargs): super().__init__( *args, alpha=0.0, dp_beta=0.1, inference_passes=1, **kwargs) if ep_beta <= 0 or not 0 < dt <= 1: raise ValueError("invalid Equilibrium Propagation scale or step") if free_steps < 1 or nudge_steps < 1: raise ValueError("EP phases require positive relaxation steps") self.ep_beta = float(ep_beta) self.ep_dt = float(dt) self.ep_free_steps = int(free_steps) self.ep_nudge_steps = int(nudge_steps) self.ep_random_beta_sign = bool(random_beta_sign) @staticmethod def ep_rho(state): return state.clamp(0.0, 1.0) @staticmethod def ep_rhop(state): return ((state >= 0.0) & (state <= 1.0)).to(state.dtype) @torch.no_grad() def ep_settle(self, image, one_hot, beta=0.0, steps=None, initial_states=None): steps = self.ep_free_steps if steps is None else int(steps) clipped_image = self.ep_rho(image) if initial_states is None: shapes = self.forward( clipped_image, training=True, update_stats=False)["hiddens"] states = [torch.zeros_like(value) for value in shapes] states.append(torch.zeros( image.shape[0], self.n_classes, device=image.device, dtype=image.dtype)) else: states = [value.detach().clone() for value in initial_states] for _ in range(steps): rho_hidden = [ self.ep_rho(value) for value in states[:-1]] fields = rho_hidden + [self.ep_rho(states[-1])] updated = [] for index, state in enumerate(states[:-1]): prediction, _ = self._node_prediction( index, rho_hidden, clipped_image) feedback = self._outgoing_feedback( index, fields, rho_hidden, clipped_image) drive = -self.ep_rho(state) + prediction + feedback updated.append(self.ep_rho( state + self.ep_dt * self.ep_rhop(state) * drive)) features = rho_hidden[-1].mean(dim=(2, 3)) prediction = features @ self.W_out.t() + self.b_out output = states[-1] drive = -self.ep_rho(output) + prediction if beta: drive = drive + 2.0 * beta * (one_hot - output) updated.append(self.ep_rho( output + self.ep_dt * self.ep_rhop(output) * drive)) states = updated return states @torch.no_grad() def _ep_phase_correlations(self, image, states): clipped_image = self.ep_rho(image) rho_hidden = [self.ep_rho(value) for value in states[:-1]] batch = image.shape[0] directions = [] gamma_directions = [] beta_directions = [] for index in range(self.n_hidden): _, cache = self._node_prediction( index, rho_hidden, clipped_image) spec = self.layer_specs[index] field = self.ep_rho(states[index]) * spec.branch_scale field, gamma_direction, beta_direction = ( self._normalization_backward( index, field, cache["normalization"])) direction = torch.nn.grad.conv2d_weight( cache["pre"], self.W[index].shape, field, stride=spec.stride, padding=spec.padding) directions.append(direction / batch) if gamma_direction is not None: gamma_directions.append(gamma_direction / batch) beta_directions.append(beta_direction / batch) features = rho_hidden[-1].mean(dim=(2, 3)) output = self.ep_rho(states[-1]) output_weight = output.t() @ features / batch output_bias = output.mean(dim=0) return ( directions, gamma_directions, beta_directions, output_weight, output_bias) @torch.no_grad() def ep_ascent_directions(self, image, free_states, nudged_states, signed_beta): if signed_beta == 0: raise ValueError("EP contrast requires a nonzero nudge") free = self._ep_phase_correlations(image, free_states) nudged = self._ep_phase_correlations(image, nudged_states) result = [] for nudged_group, free_group in zip(nudged, free): if isinstance(nudged_group, list): result.append([ (nudged_value - free_value) / signed_beta for nudged_value, free_value in zip( nudged_group, free_group)]) else: result.append((nudged_group - free_group) / signed_beta) return tuple(result) def ep_step(self, image, labels, eta, eta_output=None, momentum=0.0, weight_decay=0.0, generator=None): one_hot = F.one_hot(labels, self.n_classes).to(image.dtype) with torch.no_grad(): free_states = self.ep_settle( image, one_hot, beta=0.0, steps=self.ep_free_steps) sign = 1.0 if self.ep_random_beta_sign: draw = torch.randint( 0, 2, (), device=image.device, generator=generator) sign = 1.0 if int(draw) else -1.0 signed_beta = sign * self.ep_beta nudged_states = self.ep_settle( image, one_hot, beta=signed_beta, steps=self.ep_nudge_steps, initial_states=free_states) (directions, gamma_directions, beta_directions, output_weight, output_bias) = self.ep_ascent_directions( image, free_states, nudged_states, signed_beta) self.apply_ascent( directions, output_weight, output_bias, eta, eta_output=eta_output, momentum=momentum, weight_decay=weight_decay, gamma_directions=gamma_directions, beta_directions=beta_directions) loss = F.mse_loss(free_states[-1], one_hot) return float(loss), free_states