experiment: implement KP mixed-traffic innovation controls

1 files changed, 148 insertions, 26 deletions

diff --git a/experiments/conv_run.py b/experiments/conv_run.py
index 738c86d..d2c0653 100644
--- a/experiments/conv_run.py
+++ b/experiments/conv_run.py
@@ -12,11 +12,14 @@ import time
 import torch
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-from sdil.conv import (CIFARHierarchicalFAResNet, CIFARKPResNet,
-                       CIFARLocalResNet, CIFARSDILResNet, ConvSDILConfig,
+from sdil.conv import (CIFARHierarchicalFAResNet, CIFARKPMixedTrafficResNet,
+                       CIFARKPResNet, CIFARLocalResNet, CIFARSDILResNet,
+                       ConvSDILConfig,
                        conv_alignment_report, conv_apical_calibration_step,
                        conv_hierarchical_alignment_report,
                        conv_hierarchical_step, conv_kolen_pollack_step,
+                       conv_kp_mixed_traffic_alignment_report,
+                       conv_kp_mixed_traffic_step,
                        conv_learned_hierarchical_step, conv_local_step,
                        evaluate_conv, hierarchical_feedback_tracking_report,
                        hierarchical_parameter_subspace_calibration)
@@ -108,17 +111,25 @@ def build(args):
         bn_momentum=args.bn_momentum, bn_eps=args.bn_eps)
     if args.mode == "bp":
         return CIFARLocalResNet(**common), None
-    if args.mode in ("hfa", "lhfa", "wm", "rrm", "kp"):
-        network_class = (
-            CIFARKPResNet if args.mode == "kp" else CIFARHierarchicalFAResNet)
-        net = network_class(
-            **common, feedback_seed=args.apical_seed,
-            feedback_scale=args.a_scale)
+    if args.mode in ("hfa", "lhfa", "wm", "rrm", "kp", "kp_traffic"):
+        if args.mode == "kp_traffic":
+            net = CIFARKPMixedTrafficResNet(
+                **common, feedback_seed=args.apical_seed,
+                feedback_scale=args.a_scale, traffic_seed=args.traffic_seed)
+            net.traffic_rule = args.traffic_rule
+        else:
+            network_class = (
+                CIFARKPResNet if args.mode == "kp"
+                else CIFARHierarchicalFAResNet)
+            net = network_class(
+                **common, feedback_seed=args.apical_seed,
+                feedback_scale=args.a_scale)
         config = ConvSDILConfig(
             eta=args.lr, eta_output=args.output_lr,
             eta_A=(args.eta_A if args.mode == "lhfa" else 0.0),
+            eta_P=args.eta_P,
             momentum=args.momentum, weight_decay=args.weight_decay,
-            learn_A=args.mode == "lhfa", learn_P=False,
+            learn_A=args.mode == "lhfa", learn_P=args.mode == "kp_traffic",
             pert_sigma=args.pert_sigma, pert_every=args.pert_every,
             pert_directions=args.pert_directions,
             apical_calibration_mode="hierarchical_parameter_subspace")
@@ -145,7 +156,9 @@ def work_report(net, mode, counters):
     apical_macs = getattr(net, "apical_macs_per_example", 0)
     normal_forward = counters["ordinary_examples"] * forward_macs
     warmup_forward = ((counters["predictor_warmup_examples"]
-                       + counters["apical_warmup_examples"]) * forward_macs)
+                       + counters["apical_warmup_examples"]
+                       + counters["traffic_calibration_examples"]
+                       + counters["traffic_audit_examples"]) * forward_macs)
     calibration_forward = counters["perturbation_forward_examples"] * forward_macs
     if mode == "bp":
         bp_reverse = 2 * counters["ordinary_examples"] * forward_macs
@@ -156,7 +169,9 @@ def work_report(net, mode, counters):
         bp_reverse = 0
         local_correlation = counters["ordinary_examples"] * forward_macs
         apical_inference = ((counters["ordinary_examples"]
-                            + counters["apical_warmup_examples"]) * apical_macs)
+                            + counters["apical_warmup_examples"]
+                            + counters["traffic_calibration_examples"])
+                            * apical_macs)
         regression_multiplier = 2 if mode == "lhfa" else 1
         apical_regression = (regression_multiplier
                              * counters["calibration_event_examples"]
@@ -168,7 +183,17 @@ def work_report(net, mode, counters):
     mirror_prediction = mirror_forward if mode == "rrm" else 0
     mirror_correlation = mirror_forward
     kp_feedback_correlation = (
-        counters["ordinary_examples"] * apical_macs if mode == "kp" else 0)
+        counters["ordinary_examples"] * apical_macs
+        if mode in ("kp", "kp_traffic") else 0)
+    if mode == "kp_traffic":
+        elementwise_operations = (
+            counters["ordinary_examples"]
+            * net.mixed_elementwise_ops_per_example()
+            + (counters["predictor_warmup_examples"]
+               + counters["predictor_update_examples"])
+            * net.predictor_elementwise_ops_per_example)
+    else:
+        elementwise_operations = 0
     components = {
         "ordinary_forward_macs": normal_forward,
         "warmup_clean_forward_macs": warmup_forward,
@@ -187,12 +212,17 @@ def work_report(net, mode, counters):
         "apical_macs_per_example": apical_macs,
         "components": components,
         "total_macs_estimate": sum(components.values()),
+        "elementwise_operations_estimate": elementwise_operations,
         "total_clean_forward_examples": (
             counters["ordinary_examples"] + counters["predictor_warmup_examples"]
-            + counters["apical_warmup_examples"]),
+            + counters["apical_warmup_examples"]
+            + counters["traffic_calibration_examples"]
+            + counters["traffic_audit_examples"]),
         "total_forward_equivalent_examples": (
             counters["ordinary_examples"] + counters["predictor_warmup_examples"]
             + counters["apical_warmup_examples"]
+            + counters["traffic_calibration_examples"]
+            + counters["traffic_audit_examples"]
             + counters["perturbation_forward_examples"]),
         "logical_batch_loss_queries": counters["logical_batch_loss_queries"],
         "causal_scalar_observations": counters["causal_scalar_observations"],
@@ -203,18 +233,28 @@ def work_report(net, mode, counters):
             "multiply-accumulates in conv/linear maps; one local weight correlation "
             "equals one forward-weight MAC count; BP reverse is estimated as one "
             "weight-gradient plus one activation-gradient convolution per forward "
-            "convolution; elementwise nonlinearities and optimizer arithmetic excluded"),
+            "convolution; mixed-traffic/predictor elementwise arithmetic is reported "
+            "separately and is not folded into MACs"),
     }
 
 
 def run(args):
     if args.eval_split == "test" and args.eval_every:
         raise ValueError("test protocols must use --eval_every 0 (one final evaluation)")
-    if args.mode not in ("sdil", "lhfa") and (
+    if args.mode not in ("sdil", "lhfa", "kp_traffic") and (
             args.a_warmup_steps or args.learn_P):
         raise ValueError("apical/predictor warmup is restricted to SDIL")
     if args.mode == "lhfa" and args.learn_P:
         raise ValueError("predictor learning is not defined for learned HFA")
+    if args.mode == "kp_traffic":
+        if not args.learn_P or args.predictor_warmup_steps < 1:
+            raise ValueError("mixed-traffic KP requires predictor warmup")
+        if args.predictor_every < 1:
+            raise ValueError("mixed-traffic KP requires a predictor cadence")
+        if args.traffic_calibration_examples < 1 or args.traffic_ratio <= 0:
+            raise ValueError("invalid mixed-traffic calibration")
+    elif args.predictor_every:
+        raise ValueError("predictor cadence is restricted to mixed-traffic KP")
     if args.mode not in ("wm", "rrm") and args.mirror_warmup_steps:
         raise ValueError("mirror warmup is restricted to weight mirror mode")
     if args.mirror_every < 1 or args.mirror_batch_size < 1:
@@ -262,12 +302,17 @@ def run(args):
         "mirror_conv_examples": 0,
         "mirror_readout_examples": 0,
         "mirror_events": 0,
+        "traffic_calibration_examples": 0,
+        "traffic_audit_examples": 0,
+        "predictor_update_examples": 0,
     }
     log = {
         "schema_version": 1,
         "protocol_family": "oral_a_cifar_local_resnet_development",
         "calibration_metric_space": (
             None if config is None or args.mode == "hfa" else
+            "reciprocal_local_activity_products_with_mixed_apical_traffic"
+            if args.mode == "kp_traffic" else
             "hierarchical_feedback_parameters" if args.mode == "lhfa" else
             "reciprocal_local_activity_products" if args.mode == "kp" else {
                 "wm": "local_parent_child_response",
@@ -302,7 +347,8 @@ def run(args):
                 if args.mode == "hfa" else 0),
             "adaptive_feedback_parameters": (
                 getattr(net, "n_fixed_feedback_parameters", 0)
-                if args.mode in ("lhfa", "wm", "rrm", "kp") else 0),
+                if args.mode in ("lhfa", "wm", "rrm", "kp", "kp_traffic")
+                else 0),
         },
         "epochs": [],
     }
@@ -313,6 +359,33 @@ def run(args):
     apical_warmup_wall = 0.0
     mirror_warmup_wall = 0.0
     loader_state = train.g.get_state().clone()
+    traffic_calibration_batch = None
+    if args.mode == "kp_traffic":
+        count = min(args.traffic_calibration_examples, train.x.shape[0])
+        if count != args.traffic_calibration_examples:
+            raise ValueError("traffic calibration prefix is unavailable")
+        calibration_x = train.x[:count]
+        calibration_y = train.y[:count]
+        traffic_calibration_batch = (calibration_x, calibration_y)
+        calibration_forward = net.forward(
+            calibration_x, return_cache=True, training=True,
+            update_stats=False)
+        calibration_output_error = (
+            torch.softmax(calibration_forward["logits"], dim=1)
+            - torch.nn.functional.one_hot(
+                calibration_y, net.n_classes).to(
+                    calibration_forward["logits"].dtype))
+        calibration_instruction = net.hierarchical_teaching(
+            calibration_output_error, calibration_forward)
+        log["traffic_calibration"] = net.calibrate_traffic_gain(
+            calibration_instruction, calibration_forward["hiddens"],
+            args.traffic_ratio)
+        log["traffic_calibration"].update({
+            "examples": count,
+            "data_source": "first unaugmented development-training examples",
+            "uses_validation_endpoint": False,
+        })
+        counters["traffic_calibration_examples"] += count
     if args.mode in ("wm", "rrm") and args.mirror_warmup_steps:
         sync(args.device)
         mirror_start = time.time()
@@ -342,6 +415,7 @@ def run(args):
         sync(args.device)
         warmup_start = time.time()
         iterator = iter(train)
+        predictor_metrics = []
         for _ in range(args.predictor_warmup_steps):
             try:
                 x, _ = next(iterator)
@@ -349,10 +423,30 @@ def run(args):
                 iterator = iter(train)
                 x, _ = next(iterator)
             forward = net.forward(x, training=True, update_stats=False)
-            net.predictor_step(
-                forward["hiddens"], config.eta_P, config.nuisance_scale)
+            predictor_metric = (net.predictor_step(
+                forward["hiddens"], config.eta_P)
+                if args.mode == "kp_traffic" else net.predictor_step(
+                    forward["hiddens"], config.eta_P,
+                    config.nuisance_scale))
+            predictor_metrics.append(predictor_metric)
             counters["predictor_warmup_examples"] += x.shape[0]
         train.g.set_state(loader_state)
+        log["predictor_warmup"] = {
+            "steps": args.predictor_warmup_steps,
+            "first_mse": predictor_metrics[0],
+            "mean_mse": sum(predictor_metrics) / len(predictor_metrics),
+            "last_mse": predictor_metrics[-1],
+            "instruction_present": False,
+        }
+        if args.mode == "kp_traffic":
+            audit_x, _ = traffic_calibration_batch
+            audit_forward = net.forward(
+                audit_x, training=True, update_stats=False)
+            log["predictor_warmup"][
+                "post_warmup_traffic_residual_rms_ratio"] = (
+                    net.predictor_traffic_residual_rms_ratio(
+                        audit_forward["hiddens"]))
+            counters["traffic_audit_examples"] += audit_x.shape[0]
         sync(args.device)
         predictor_warmup_wall = time.time() - warmup_start
 
@@ -426,6 +520,7 @@ def run(args):
         examples = 0
         calibration_metrics = []
         mirror_metrics = []
+        signal_metrics = []
         for x, y in train:
             batch = x.shape[0]
             if args.mode == "bp":
@@ -476,6 +571,17 @@ def run(args):
                 result = conv_kolen_pollack_step(net, x, y, config)
                 loss = result["loss"]
                 did_perturb = False
+            elif args.mode == "kp_traffic":
+                result = conv_kp_mixed_traffic_step(
+                    net, x, y, config, step, args.traffic_rule,
+                    args.predictor_every)
+                loss = result["loss"]
+                did_perturb = False
+                signal_metrics.append({key: result[key] for key in (
+                    "teaching_rms", "instruction_rms", "raw_apical_rms",
+                    "innovation_rms", "traffic_rms")})
+                if result["did_predictor_update"]:
+                    counters["predictor_update_examples"] += batch
             else:
                 result = conv_local_step(
                     net, x, y, config, step, generator=perturb_generator)
@@ -521,7 +627,12 @@ def run(args):
                 key: sum(value[key] for value in mirror_metrics)
                 / len(mirror_metrics) for key in mirror_metrics[0]
             }
-        if args.mode == "kp":
+        if signal_metrics:
+            record["mixed_apical"] = {
+                key: sum(value[key] for value in signal_metrics)
+                / len(signal_metrics) for key in signal_metrics[0]
+            }
+        if args.mode in ("kp", "kp_traffic"):
             tracking = hierarchical_feedback_tracking_report(net)
             record["feedback_tracking"] = {
                 "mean_feedback_forward_cosine": tracking[
@@ -569,11 +680,15 @@ def run(args):
         probe = min(args.alignment_probe, train.x.shape[0])
         sync(args.device)
         diagnostic_start = time.time()
-        diagnostics = (conv_hierarchical_alignment_report(
-            net, train.x[:probe], train.y[:probe])
-            if args.mode in ("hfa", "lhfa", "wm", "rrm", "kp") else
-            conv_alignment_report(
-                net, train.x[:probe], train.y[:probe], config))
+        if args.mode == "kp_traffic":
+            diagnostics = conv_kp_mixed_traffic_alignment_report(
+                net, train.x[:probe], train.y[:probe], args.traffic_rule)
+        elif args.mode in ("hfa", "lhfa", "wm", "rrm", "kp"):
+            diagnostics = conv_hierarchical_alignment_report(
+                net, train.x[:probe], train.y[:probe])
+        else:
+            diagnostics = conv_alignment_report(
+                net, train.x[:probe], train.y[:probe], config)
         sync(args.device)
         diagnostics["wall_s"] = time.time() - diagnostic_start
         values = diagnostics["teaching_negative_gradient_cosine"]
@@ -620,8 +735,8 @@ def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--mode", choices=(
-            "bp", "dfa", "hfa", "lhfa", "wm", "rrm", "kp", "sdil",
-            "nodepert"),
+            "bp", "dfa", "hfa", "lhfa", "wm", "rrm", "kp",
+            "kp_traffic", "sdil", "nodepert"),
         required=True)
     parser.add_argument("--out", required=True)
     parser.add_argument("--device", default="cpu")
@@ -673,6 +788,13 @@ def parse_args():
         choices=("unit_targets", "channel_subspace", "vectorizer_subspace"),
         default="unit_targets")
     parser.add_argument("--predictor_warmup_steps", type=int, default=0)
+    parser.add_argument("--predictor_every", type=int, default=0)
+    parser.add_argument("--traffic_rule",
+                        choices=("raw", "matched", "innovation"),
+                        default="innovation")
+    parser.add_argument("--traffic_seed", type=int, default=4000)
+    parser.add_argument("--traffic_ratio", type=float, default=4.0)
+    parser.add_argument("--traffic_calibration_examples", type=int, default=64)
     parser.add_argument("--a_warmup_steps", type=int, default=0)
     parser.add_argument("--mirror_warmup_steps", type=int, default=0)
     parser.add_argument("--mirror_every", type=int, default=16)