oral-a: make BatchNorm credit local and causal

2 files changed, 150 insertions, 11 deletions

diff --git a/experiments/conv_local_smoke.py b/experiments/conv_local_smoke.py
index 4a0f47d..d2428eb 100644
--- a/experiments/conv_local_smoke.py
+++ b/experiments/conv_local_smoke.py
@@ -24,6 +24,11 @@ def architecture_checks():
         assert net.n_forward_parameters == parameters
         assert len(net.blocks) == 3 * ((depth - 2) // 6)
         assert len(net.W) + 1 == depth
+    batchnorm_parameters = {8: 75290, 20: 269722, 32: 464154, 56: 853018}
+    for depth, parameters in batchnorm_parameters.items():
+        net = CIFARLocalResNet(
+            depth=depth, normalization="batchnorm", residual_scale=1.0)
+        assert net.n_forward_parameters == parameters
     for depth in (7, 9, 21):
         try:
             CIFARLocalResNet(depth=depth)
@@ -62,8 +67,10 @@ def exact_local_gradient_check():
     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(
+    (directions, gamma_directions, beta_directions,
+     out_direction, bias_direction) = local.local_ascent_directions(
         teaching, output_error, forward)
+    assert gamma_directions == beta_directions == []
 
     relative_errors = []
     for direction, parameter in zip(directions, local.W):
@@ -94,6 +101,78 @@ def exact_local_gradient_check():
     }
 
 
+def exact_batchnorm_local_gradient_check():
+    torch.manual_seed(31)
+    batch = 4
+    x = torch.randn(batch, 3, 32, 32)
+    y = torch.tensor([0, 1, 2, 3])
+    common = dict(
+        depth=8, base_width=2, seed=29,
+        normalization="batchnorm", residual_scale=1.0)
+    local = CIFARLocalResNet(**common)
+    bp = CIFARLocalResNet(**common)
+    parameters = local.W + local.gamma + local.beta + [local.W_out, local.b_out]
+    for parameter in parameters:
+        parameter.requires_grad_(True)
+    forward = local.forward(
+        x, return_cache=True, training=True, update_stats=True)
+    for hidden in forward["hiddens"]:
+        hidden.retain_grad()
+    loss = F.cross_entropy(forward["logits"], y)
+    loss.backward()
+    teaching = [-batch * hidden.grad for hidden in forward["hiddens"]]
+    output_error = (torch.softmax(forward["logits"].detach(), dim=1)
+                    - F.one_hot(y, 10))
+    (directions, gamma_directions, beta_directions,
+     out_direction, bias_direction) = local.local_ascent_directions(
+        teaching, output_error, forward)
+    groups = (
+        (directions, local.W),
+        (gamma_directions, local.gamma),
+        (beta_directions, local.beta),
+    )
+    relative_errors = []
+    for direction_group, parameter_group in groups:
+        for direction, parameter in zip(direction_group, parameter_group):
+            absolute = (direction + parameter.grad).abs().max()
+            relative_errors.append(float(
+                absolute / parameter.grad.abs().max().clamp_min(1e-12)))
+    assert max(relative_errors) < 3e-5
+    for parameter in parameters:
+        parameter.requires_grad_(False)
+
+    eta = 0.013
+    local.apply_ascent(
+        directions, out_direction, bias_direction, eta,
+        gamma_directions=gamma_directions, beta_directions=beta_directions)
+    bp.bp_step(x, y, eta)
+    parameter_differences = [
+        float((left - right).abs().max())
+        for left, right 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_differences = [
+        float((left - right).abs().max())
+        for left, right in zip(
+            local.running_mean + local.running_var,
+            bp.running_mean + bp.running_var)]
+    assert max(parameter_differences) < 2e-7
+    assert max(running_differences) == 0.0
+
+    running_before = [value.clone() for value in local.running_mean + local.running_var]
+    clean = local.forward(x, training=True, update_stats=False)
+    simultaneous_conv_node_perturbation(
+        local, x, y, clean, sigma=1e-3, n_directions=1,
+        generator=torch.Generator(device="cpu").manual_seed(9))
+    assert all(torch.equal(before, after) for before, after in zip(
+        running_before, local.running_mean + local.running_var))
+    assert torch.equal(local.logits(x), local.logits(x))
+    return {
+        "batchnorm_max_relative_local_gradient_error": max(relative_errors),
+        "batchnorm_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)
@@ -132,7 +211,9 @@ def perturbation_estimator_check():
         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]
+    derivative = diagnostics["directional_derivatives"][0]
+    assert derivative["coupling"] == "per_example_objective"
+    finite_difference = derivative["scaled_directional"]
     exact = sum(
         (batch * hidden.grad * direction).flatten(1).sum(dim=1)
         for hidden, direction in zip(clean["hiddens"], directions))
@@ -143,7 +224,46 @@ def perturbation_estimator_check():
         assert torch.equal(target, expected)
     for parameter in parameters:
         parameter.requires_grad_(False)
-    return float(relative.max())
+    no_norm_relative = float(relative.max())
+
+    batch = 3
+    batchnorm = CIFARSDILResNet(
+        depth=8, base_width=2, seed=14, dtype=torch.float64,
+        normalization="batchnorm", residual_scale=1.0)
+    xb = torch.randn(batch, 3, 32, 32, dtype=torch.float64)
+    yb = torch.tensor([1, 4, 9])
+    parameters = (batchnorm.W + batchnorm.gamma + batchnorm.beta
+                  + [batchnorm.W_out, batchnorm.b_out])
+    for parameter in parameters:
+        parameter.requires_grad_(True)
+    clean = batchnorm.forward(xb, training=True, update_stats=False)
+    for hidden in clean["hiddens"]:
+        hidden.retain_grad()
+    F.cross_entropy(clean["logits"], yb).backward()
+    targets, diagnostics = simultaneous_conv_node_perturbation(
+        batchnorm, xb, yb, clean, sigma=1e-6, n_directions=1,
+        generator=torch.Generator(device="cpu").manual_seed(101),
+        return_diagnostics=True)
+    directions = diagnostics["directions"][0]
+    derivative = diagnostics["directional_derivatives"][0]
+    assert derivative["coupling"] == "batch_objective"
+    scaled = derivative["scaled_directional"]
+    exact_sum_directional = sum(
+        float((batch * hidden.grad * direction).sum())
+        for hidden, direction in zip(clean["hiddens"], directions))
+    batchnorm_relative = abs(float(scaled[0]) - exact_sum_directional) / max(
+        abs(exact_sum_directional), 1e-12)
+    assert batchnorm_relative < 2e-6
+    assert torch.equal(scaled, scaled[:1].expand_as(scaled))
+    for target, direction in zip(targets, directions):
+        expected = -scaled[:, None, None, None] * direction
+        assert torch.equal(target, expected)
+    for parameter in parameters:
+        parameter.requires_grad_(False)
+    return {
+        "perturbation_jvp_max_relative_error": no_norm_relative,
+        "batchnorm_batch_objective_jvp_relative_error": batchnorm_relative,
+    }
 
 
 def apical_learning_checks():
@@ -197,7 +317,8 @@ def main():
     architecture_checks()
     perturbation_checks()
     report = exact_local_gradient_check()
-    report["perturbation_jvp_max_relative_error"] = perturbation_estimator_check()
+    report.update(exact_batchnorm_local_gradient_check())
+    report.update(perturbation_estimator_check())
     report.update(apical_learning_checks())
     print(report)
     print("ALL CONVOLUTIONAL LOCAL-ELIGIBILITY CHECKS PASSED")
diff --git a/experiments/conv_run.py b/experiments/conv_run.py
index 48eebd1..dd3f6c5 100644
--- a/experiments/conv_run.py
+++ b/experiments/conv_run.py
@@ -91,10 +91,14 @@ def scheduled_lr(base, epoch, args):
 
 
 def build(args):
+    residual_scale = args.residual_scale
+    if residual_scale is None and args.normalization == "batchnorm":
+        residual_scale = 1.0
     common = dict(
         depth=args.depth, base_width=args.width, n_classes=10,
         device=args.device, seed=args.seed, weight_scale=args.weight_scale,
-        residual_scale=args.residual_scale)
+        residual_scale=residual_scale, normalization=args.normalization,
+        bn_momentum=args.bn_momentum, bn_eps=args.bn_eps)
     if args.mode == "bp":
         return CIFARLocalResNet(**common), None
     net = CIFARSDILResNet(
@@ -151,7 +155,8 @@ def work_report(net, mode, counters):
             + counters["apical_warmup_examples"]
             + counters["perturbation_forward_examples"]),
         "logical_batch_loss_queries": counters["logical_batch_loss_queries"],
-        "scalar_loss_evaluations": counters["scalar_loss_evaluations"],
+        "causal_scalar_observations": counters["causal_scalar_observations"],
+        "per_example_cross_entropy_terms": counters["per_example_loss_terms"],
         "definition": (
             "multiply-accumulates in conv/linear maps; one local weight correlation "
             "equals one forward-weight MAC count; BP reverse is estimated as one "
@@ -198,7 +203,8 @@ def run(args):
         "perturbation_forward_examples": 0,
         "calibration_event_examples": 0,
         "logical_batch_loss_queries": 0,
-        "scalar_loss_evaluations": 0,
+        "causal_scalar_observations": 0,
+        "per_example_loss_terms": 0,
         "perturbation_events": 0,
     }
     log = {
@@ -208,7 +214,10 @@ def run(args):
         "provenance": provenance(),
         "split": split,
         "architecture": {
-            "family": "normalization-free CIFAR 6n+2 ResNet, option-A shortcuts",
+            "family": "CIFAR 6n+2 ResNet, option-A shortcuts",
+            "normalization": net.normalization,
+            "bn_momentum": net.bn_momentum if net.normalization == "batchnorm" else None,
+            "bn_eps": net.bn_eps if net.normalization == "batchnorm" else None,
             "depth": net.depth,
             "blocks_per_stage": net.blocks_per_stage,
             "base_width": net.base_width,
@@ -231,7 +240,7 @@ def run(args):
             except StopIteration:
                 iterator = iter(train)
                 x, _ = next(iterator)
-            forward = net.forward(x)
+            forward = net.forward(x, training=True, update_stats=False)
             net.predictor_step(
                 forward["hiddens"], config.eta_P, config.nuisance_scale)
             counters["predictor_warmup_examples"] += x.shape[0]
@@ -255,7 +264,9 @@ def run(args):
                 2 * config.pert_directions * batch)
             counters["calibration_event_examples"] += batch
             counters["logical_batch_loss_queries"] += 2 * config.pert_directions
-            counters["scalar_loss_evaluations"] += 2 * config.pert_directions * batch
+            counters["causal_scalar_observations"] += 2 * config.pert_directions * (
+                1 if net.normalization == "batchnorm" else batch)
+            counters["per_example_loss_terms"] += 2 * config.pert_directions * batch
             counters["perturbation_events"] += 1
         train.g.set_state(loader_state)
         log["apical_warmup"] = {
@@ -304,7 +315,10 @@ def run(args):
                     2 * config.pert_directions * batch)
                 counters["calibration_event_examples"] += batch
                 counters["logical_batch_loss_queries"] += 2 * config.pert_directions
-                counters["scalar_loss_evaluations"] += (
+                counters["causal_scalar_observations"] += (
+                    2 * config.pert_directions
+                    * (1 if net.normalization == "batchnorm" else batch))
+                counters["per_example_loss_terms"] += (
                     2 * config.pert_directions * batch)
                 counters["perturbation_events"] += 1
             step += 1
@@ -435,6 +449,10 @@ def parse_args():
     parser.add_argument("--weight_decay", type=float, default=5e-4)
     parser.add_argument("--weight_scale", type=float, default=1.0)
     parser.add_argument("--residual_scale", type=float)
+    parser.add_argument("--normalization", choices=("batchnorm", "none"),
+                        default="batchnorm")
+    parser.add_argument("--bn_momentum", type=float, default=0.1)
+    parser.add_argument("--bn_eps", type=float, default=1e-5)
     parser.add_argument("--a_scale", type=float, default=1.0)
     parser.add_argument("--eta_A", type=float, default=0.01)
     parser.add_argument("--eta_P", type=float, default=0.01)