oral-a: make BatchNorm credit local and causal

1 files changed, 25 insertions, 7 deletions

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)