diff options
Diffstat (limited to 'sdil/conv_crossover.py')
| -rw-r--r-- | sdil/conv_crossover.py | 91 |
1 files changed, 91 insertions, 0 deletions
diff --git a/sdil/conv_crossover.py b/sdil/conv_crossover.py index d455f21..329ad3b 100644 --- a/sdil/conv_crossover.py +++ b/sdil/conv_crossover.py @@ -11,6 +11,97 @@ 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 CIFARDualPropResNet(CIFARLocalResNet): """Dual Propagation on the residual DAG with author DP-transpose updates. |