diff options
| -rw-r--r-- | experiments/resnet_crossover_smoke.py | 89 | ||||
| -rw-r--r-- | sdil/conv_crossover.py | 91 |
2 files changed, 178 insertions, 2 deletions
diff --git a/experiments/resnet_crossover_smoke.py b/experiments/resnet_crossover_smoke.py index fdc92fa..aadede5 100644 --- a/experiments/resnet_crossover_smoke.py +++ b/experiments/resnet_crossover_smoke.py @@ -8,7 +8,7 @@ import torch sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from sdil.conv import CIFARLocalResNet -from sdil.conv_crossover import CIFARDualPropResNet +from sdil.conv_crossover import CIFARDualPropResNet, CIFARPEPITAResNet def relative_error(actual, expected): @@ -85,8 +85,93 @@ def audit_dualprop(): } +def audit_pepita(): + common = dict( + depth=8, base_width=2, seed=81, normalization="batchnorm", + residual_scale=1.0, dtype=torch.float64) + reference = CIFARLocalResNet(**common) + net = CIFARPEPITAResNet( + **common, projection_scale=0.05, projection_seed=82) + generator = torch.Generator().manual_seed(83) + image = torch.randn(3, 3, 32, 32, generator=generator, + dtype=torch.float64) + labels = torch.tensor([1, 4, 8]) + one_hot = torch.nn.functional.one_hot(labels, 10).to(torch.float64) + with torch.no_grad(): + reference_output = reference.forward( + image, training=True, update_stats=False)["logits"] + clean = net.forward( + image, return_cache=True, training=True, update_stats=False) + clean_error = torch.softmax(clean["logits"], dim=1) - one_hot + input_error = torch.einsum( + "bc,cijk->bijk", clean_error, net.input_feedback) + modulated = net.forward( + image + input_error, return_cache=True, + training=True, update_stats=False) + modulated_error = ( + torch.softmax(modulated["logits"], dim=1) - one_hot) + forward_error = float(torch.max(torch.abs( + reference_output - clean["logits"]))) + + parameters = ( + net.W + net.gamma + net.beta + [net.W_out, net.b_out]) + for parameter in parameters: + parameter.requires_grad_(True) + objective = image.new_zeros(()) + batch = image.shape[0] + for index, (clean_hidden, modulated_hidden, cache, spec) in enumerate(zip( + clean["hiddens"], modulated["hiddens"], modulated["caches"], + net.layer_specs)): + field = (clean_hidden - modulated_hidden).detach() + convolution = torch.nn.functional.conv2d( + cache["pre"].detach(), net.W[index], + stride=spec.stride, padding=spec.padding) + normalized, _ = net._normalize( + index, convolution, training=True, update_stats=False) + prediction = spec.branch_scale * normalized + spatial = prediction.shape[2] * prediction.shape[3] + objective += torch.sum(field * prediction) / (batch * spatial) + output_prediction = ( + modulated["features"].detach() @ net.W_out.t() + net.b_out) + objective += torch.sum( + modulated_error.detach() * output_prediction) / batch + gradients = torch.autograd.grad(objective, parameters) + (directions, gamma_directions, beta_directions, + output_weight, output_bias) = net.pepita_ascent_directions( + clean, modulated, modulated_error) + actual = ( + directions + gamma_directions + beta_directions + + [output_weight, output_bias]) + errors = [ + relative_error(direction, -gradient) + for direction, gradient in zip(actual, gradients)] + for parameter in parameters: + parameter.requires_grad_(False) + projection_limit = (6.0 / (3 * 32 * 32)) ** 0.5 * 0.05 + observed_limit = float(torch.max(torch.abs(net.input_feedback))) + if (forward_error >= 1e-12 or max(errors) >= 2e-12 + or observed_limit > projection_limit): + raise AssertionError({ + "forward_error": forward_error, + "direction_errors": errors, + "projection_limit": projection_limit, + "observed_limit": observed_limit, + }) + return { + "matched_forward_max_absolute_error": forward_error, + "local_equation_max_relative_error": max(errors), + "input_projection_shape": list(net.input_feedback.shape), + "input_projection_limit": projection_limit, + "observed_input_projection_max": observed_limit, + "uses_reverse_task_loss_graph": False, + } + + def main(): - print(json.dumps({"dualprop": audit_dualprop()}, + print(json.dumps({ + "dualprop": audit_dualprop(), + "pepita": audit_pepita(), + }, indent=2, sort_keys=True)) 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. |