#!/usr/bin/env python3 """Prove the convolutional local eligibility matches exact BP when instructed.""" import os import sys import torch import torch.nn.functional as F sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from sdil.conv import (CIFARLocalResNet, CIFARSDILResNet, ConvSDILConfig, conv_local_step, simultaneous_conv_node_perturbation) def architecture_checks(): expected = { 8: (7, 74810), 20: (19, 268346), 32: (31, 461882), 56: (55, 848954), } for depth, (hidden, parameters) in expected.items(): net = CIFARLocalResNet(depth=depth) assert net.n_hidden == hidden assert net.n_forward_parameters == parameters assert len(net.blocks) == 3 * ((depth - 2) // 6) assert len(net.W) + 1 == depth for depth in (7, 9, 21): try: CIFARLocalResNet(depth=depth) except ValueError: pass else: raise AssertionError(f"invalid depth {depth} was accepted") x = torch.arange(2 * 3 * 8 * 8, dtype=torch.float32).reshape(2, 3, 8, 8) shortcut = CIFARLocalResNet._option_a_shortcut(x, 6, 2) assert tuple(shortcut.shape) == (2, 6, 4, 4) assert torch.count_nonzero(shortcut[:, 0]) == 0 assert torch.count_nonzero(shortcut[:, -2:]) == 0 assert torch.equal(shortcut[:, 1:4], x[:, :, ::2, ::2]) def exact_local_gradient_check(): torch.manual_seed(123) batch = 3 x = torch.randn(batch, 3, 32, 32) y = torch.tensor([0, 3, 8]) local = CIFARLocalResNet(depth=8, base_width=4, seed=19) bp = CIFARLocalResNet(depth=8, base_width=4, seed=19) parameters = local.W + [local.W_out, local.b_out] for parameter in parameters: parameter.requires_grad_(True) forward = local.forward(x, return_cache=True) for hidden in forward["hiddens"]: hidden.retain_grad() loss = F.cross_entropy(forward["logits"], y) loss.backward() # .backward() differentiated a batch-mean loss. Multiplying hidden grads # by B recovers the per-example convention consumed by the local rule. teaching = [-batch * hidden.grad for hidden in forward["hiddens"]] output_error = (torch.softmax(forward["logits"].detach(), dim=1) - F.one_hot(y, local.n_classes)) directions, out_direction, bias_direction = local.local_ascent_directions( teaching, output_error, forward) relative_errors = [] for direction, parameter in zip(directions, local.W): absolute = (direction + parameter.grad).abs().max() scale = parameter.grad.abs().max().clamp_min(1e-12) relative_errors.append(float(absolute / scale)) output_abs = float((out_direction + local.W_out.grad).abs().max()) bias_abs = float((bias_direction + local.b_out.grad).abs().max()) assert max(relative_errors) < 3e-5 assert output_abs < 2e-6 and bias_abs < 2e-6 for parameter in parameters: parameter.requires_grad_(False) eta = 0.017 local.apply_ascent(directions, out_direction, bias_direction, eta) bp_loss = bp.bp_step(x, y, eta) assert abs(float(loss.detach()) - bp_loss) < 1e-7 parameter_differences = [ float((left - right).abs().max()) for left, right in zip( local.W + [local.W_out, local.b_out], bp.W + [bp.W_out, bp.b_out])] assert max(parameter_differences) < 2e-7 return { "max_relative_local_gradient_error": max(relative_errors), "output_absolute_error": output_abs, "post_update_parameter_max_error": max(parameter_differences), } def perturbation_checks(): net = CIFARLocalResNet(depth=8, base_width=4, seed=3) x = torch.randn(2, 3, 32, 32) clean = net.forward(x) perturbations = [torch.zeros_like(hidden) for hidden in clean["hiddens"]] perturbed = net.forward(x, perturbations=perturbations) assert torch.equal(clean["logits"], perturbed["logits"]) perturbations[0] = torch.ones_like(perturbations[0]) * 0.01 changed = net.forward(x, perturbations=perturbations) assert not torch.equal(clean["logits"], changed["logits"]) try: net.forward(x, perturbations=perturbations[:-1]) except ValueError: pass else: raise AssertionError("short perturbation list was accepted") def perturbation_estimator_check(): """Antithetic finite differences equal the simultaneous hidden JVP.""" torch.manual_seed(3) batch = 2 net = CIFARSDILResNet( depth=8, base_width=2, seed=4, dtype=torch.float64) x = torch.randn(batch, 3, 32, 32, dtype=torch.float64) y = torch.tensor([2, 7]) parameters = net.W + [net.W_out, net.b_out] for parameter in parameters: parameter.requires_grad_(True) clean = net.forward(x, return_cache=True) for hidden in clean["hiddens"]: hidden.retain_grad() F.cross_entropy(clean["logits"], y).backward() generator = torch.Generator(device="cpu").manual_seed(99) targets, diagnostics = simultaneous_conv_node_perturbation( net, x, y, clean, sigma=1e-6, n_directions=1, generator=generator, return_diagnostics=True) directions = diagnostics["directions"][0] finite_difference = diagnostics["directional_derivatives"][0] exact = sum( (batch * hidden.grad * direction).flatten(1).sum(dim=1) for hidden, direction in zip(clean["hiddens"], directions)) relative = (finite_difference - exact).abs() / exact.abs().clamp_min(1e-12) assert float(relative.max()) < 2e-6 for target, direction in zip(targets, directions): expected = -finite_difference[:, None, None, None] * direction assert torch.equal(target, expected) for parameter in parameters: parameter.requires_grad_(False) return float(relative.max()) def apical_learning_checks(): torch.manual_seed(11) net = CIFARSDILResNet(depth=8, base_width=2, seed=6) x = torch.randn(8, 3, 32, 32) y = torch.arange(8) % 10 clean = net.forward(x, return_cache=True) output_signal = (torch.softmax(clean["logits"], dim=1) - F.one_hot(y, 10)) prediction, _, _ = net.apical_components( output_signal, clean["hiddens"], use_residual=True) targets = [torch.randn_like(value) * 0.01 for value in prediction] before = sum(float((target - value).square().sum()) for target, value in zip(targets, prediction)) net.calibrate_apical(output_signal, prediction, targets, eta=0.1) after_prediction, _, _ = net.apical_components( output_signal, clean["hiddens"], use_residual=True) after = sum(float((target - value).square().sum()) for target, value in zip(targets, after_prediction)) assert after < before predictor_net = CIFARSDILResNet(depth=8, base_width=2, seed=8) hiddens = [torch.randn(64, *shape) for shape in predictor_net.hidden_shapes] initial = predictor_net.predictor_step(hiddens, eta=0.1, nuisance_scale=0.5) final = initial for _ in range(60): final = predictor_net.predictor_step(hiddens, eta=0.1, nuisance_scale=0.5) assert final < initial * 1e-3 weights_before = [weight.clone() for weight in net.W] result = conv_local_step( net, x[:2], y[:2], ConvSDILConfig( eta=1e-3, eta_A=1e-3, momentum=0.0, weight_decay=0.0, pert_every=1), step=0, generator=torch.Generator(device="cpu").manual_seed(7)) assert result["did_perturb"] and result["calibration"] is not None assert torch.isfinite(torch.tensor(list( value for key, value in result.items() if isinstance(value, float) and key != "predictor_mse"))).all() assert any(not torch.equal(before_weight, after_weight) for before_weight, after_weight in zip(weights_before, net.W)) assert all(not parameter.requires_grad for parameter in net.W + [net.W_out, net.b_out]) return {"apical_mse_ratio": after / before, "predictor_mse_ratio": final / initial} def main(): architecture_checks() perturbation_checks() report = exact_local_gradient_check() report["perturbation_jvp_max_relative_error"] = perturbation_estimator_check() report.update(apical_learning_checks()) print(report) print("ALL CONVOLUTIONAL LOCAL-ELIGIBILITY CHECKS PASSED") if __name__ == "__main__": main()