diff options
Diffstat (limited to 'experiments/conv_local_smoke.py')
| -rw-r--r-- | experiments/conv_local_smoke.py | 81 |
1 files changed, 79 insertions, 2 deletions
diff --git a/experiments/conv_local_smoke.py b/experiments/conv_local_smoke.py index 8070d9f..9207e04 100644 --- a/experiments/conv_local_smoke.py +++ b/experiments/conv_local_smoke.py @@ -7,8 +7,10 @@ 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, - channel_subspace_apical_calibration, conv_local_step, +from sdil.conv import (CIFARHierarchicalFAResNet, CIFARLocalResNet, + CIFARSDILResNet, ConvSDILConfig, + channel_subspace_apical_calibration, + conv_hierarchical_step, conv_local_step, simultaneous_conv_node_perturbation, vectorizer_subspace_apical_calibration) @@ -597,6 +599,80 @@ def vectorizer_subspace_estimator_check(): } +def hierarchical_feedback_checks(): + """The residual feedback graph becomes exact only under an audit copy.""" + torch.manual_seed(91) + exact = CIFARHierarchicalFAResNet( + depth=8, base_width=2, seed=92, dtype=torch.float64, + normalization="batchnorm", residual_scale=1.0) + exact.Q = [value.clone() for value in exact.W] + exact.R_out.copy_(-exact.W_out.t()) + x = torch.randn(3, 3, 32, 32, dtype=torch.float64) + y = torch.tensor([1, 5, 8]) + parameters = (exact.W + exact.gamma + exact.beta + + [exact.W_out, exact.b_out]) + for parameter in parameters: + parameter.requires_grad_(True) + forward = exact.forward(x, return_cache=True, training=True, + update_stats=False) + gradients = torch.autograd.grad( + F.cross_entropy(forward["logits"], y), forward["hiddens"]) + output_error = (torch.softmax(forward["logits"].detach(), dim=1) + - F.one_hot(y, 10).to(torch.float64)) + teaching = exact.hierarchical_teaching(output_error, forward) + relative = [float((signal + x.shape[0] * gradient).abs().max() + / gradient.abs().max().clamp_min(1e-30) + / x.shape[0]) + for signal, gradient in zip(teaching, gradients)] + assert max(relative) < 2e-12 + for parameter in parameters: + parameter.requires_grad_(False) + + # With independent feedback the forward model is bitwise unchanged, while + # changing only the feedback seed changes Q/R. This is the actual baseline. + left = CIFARHierarchicalFAResNet( + depth=8, base_width=2, seed=93, feedback_seed=1001) + right = CIFARHierarchicalFAResNet( + depth=8, base_width=2, seed=93, feedback_seed=1002) + assert all(torch.equal(a, b) for a, b in zip( + left.W + [left.W_out], right.W + [right.W_out])) + assert any(not torch.equal(a, b) for a, b in zip( + left.Q + [left.R_out], right.Q + [right.R_out])) + + # The local update must match exact BP when feedback is explicitly copied + # in this audit-only comparator. Actual HFA never performs this copy. + local = CIFARHierarchicalFAResNet( + depth=8, base_width=2, seed=94, normalization="batchnorm", + residual_scale=1.0) + bp = CIFARLocalResNet( + depth=8, base_width=2, seed=94, normalization="batchnorm", + residual_scale=1.0) + local.Q = [value.clone() for value in local.W] + local.R_out.copy_(-local.W_out.t()) + xf = torch.randn(4, 3, 32, 32) + yf = torch.tensor([0, 2, 5, 9]) + eta = 0.013 + conv_hierarchical_step( + local, xf, yf, ConvSDILConfig( + eta=eta, eta_output=eta, eta_A=0.0, momentum=0.0, + weight_decay=0.0, learn_A=False)) + bp.bp_step(xf, yf, eta, momentum=0.0, weight_decay=0.0) + parameter_error = max(float((a - b).abs().max()) for a, b in zip( + local.W + local.gamma + local.beta + [local.W_out, local.b_out], + bp.W + bp.gamma + bp.beta + [bp.W_out, bp.b_out])) + running_error = max(float((a - b).abs().max()) for a, b in zip( + local.running_mean + local.running_var, + bp.running_mean + bp.running_var)) + assert parameter_error < 2e-7 + assert running_error == 0.0 + return { + "hierarchical_symmetric_hidden_relative_error": max(relative), + "hierarchical_symmetric_update_absolute_error": parameter_error, + "hierarchical_feedback_to_forward_mac_ratio": ( + local.apical_macs_per_example / local.forward_macs_per_example), + } + + def apical_learning_checks(): torch.manual_seed(11) net = CIFARSDILResNet(depth=8, base_width=2, seed=6) @@ -702,6 +778,7 @@ def main(): report.update(perturbation_estimator_check()) report.update(channel_subspace_estimator_check()) report.update(vectorizer_subspace_estimator_check()) + report.update(hierarchical_feedback_checks()) report.update(apical_learning_checks()) print(report) print("ALL CONVOLUTIONAL LOCAL-ELIGIBILITY CHECKS PASSED") |