baseline: add convolutional hierarchical feedback alignment

2 files changed, 108 insertions, 8 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")
diff --git a/experiments/conv_run.py b/experiments/conv_run.py
index 3c0de00..162a8be 100644
--- a/experiments/conv_run.py
+++ b/experiments/conv_run.py
@@ -12,9 +12,11 @@ import time
 import torch
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-from sdil.conv import (CIFARLocalResNet, CIFARSDILResNet, ConvSDILConfig,
+from sdil.conv import (CIFARHierarchicalFAResNet, CIFARLocalResNet,
+                       CIFARSDILResNet, ConvSDILConfig,
                        conv_alignment_report, conv_apical_calibration_step,
-                       conv_local_step, evaluate_conv)
+                       conv_hierarchical_alignment_report,
+                       conv_hierarchical_step, conv_local_step, evaluate_conv)
 from sdil.data import DATA_DIR, get_cifar_image_splits
 
 
@@ -101,6 +103,16 @@ def build(args):
         bn_momentum=args.bn_momentum, bn_eps=args.bn_eps)
     if args.mode == "bp":
         return CIFARLocalResNet(**common), None
+    if args.mode == "hfa":
+        net = CIFARHierarchicalFAResNet(
+            **common, feedback_seed=args.apical_seed,
+            feedback_scale=args.a_scale)
+        config = ConvSDILConfig(
+            eta=args.lr, eta_output=args.output_lr, eta_A=0.0,
+            momentum=args.momentum, weight_decay=args.weight_decay,
+            learn_A=False, learn_P=False)
+        config.validate()
+        return net, config
     net = CIFARSDILResNet(
         **common, a_scale=args.a_scale, apical_seed=args.apical_seed,
         vectorizer_mode=args.vectorizer_mode)
@@ -213,7 +225,7 @@ def run(args):
         "schema_version": 1,
         "protocol_family": "oral_a_cifar_local_resnet_development",
         "calibration_metric_space": (
-            None if config is None else {
+            None if config is None or args.mode == "hfa" else {
                 "unit_targets": "full_hidden_field",
                 "channel_subspace": "channel_basis_moments",
                 "vectorizer_subspace": "vectorizer_parameter_gradients",
@@ -238,6 +250,8 @@ def run(args):
             "vectorizer_mode": getattr(net, "vectorizer_mode", None),
             "fixed_traffic_coefficients": getattr(
                 net, "n_fixed_traffic_coefficients", 0),
+            "fixed_feedback_parameters": getattr(
+                net, "n_fixed_feedback_parameters", 0),
         },
         "epochs": [],
     }
@@ -328,6 +342,10 @@ def run(args):
                     x, y, lr, momentum=args.momentum,
                     weight_decay=args.weight_decay)
                 did_perturb = False
+            elif args.mode == "hfa":
+                result = conv_hierarchical_step(net, x, y, config)
+                loss = result["loss"]
+                did_perturb = False
             else:
                 result = conv_local_step(
                     net, x, y, config, step, generator=perturb_generator)
@@ -404,8 +422,11 @@ def run(args):
         probe = min(args.alignment_probe, train.x.shape[0])
         sync(args.device)
         diagnostic_start = time.time()
-        diagnostics = conv_alignment_report(
-            net, train.x[:probe], train.y[:probe], config)
+        diagnostics = (conv_hierarchical_alignment_report(
+            net, train.x[:probe], train.y[:probe])
+            if args.mode == "hfa" else
+            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"]
@@ -449,7 +470,9 @@ def run(args):
 
 def parse_args():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--mode", choices=("bp", "dfa", "sdil", "nodepert"), required=True)
+    parser.add_argument(
+        "--mode", choices=("bp", "dfa", "hfa", "sdil", "nodepert"),
+        required=True)
     parser.add_argument("--out", required=True)
     parser.add_argument("--device", default="cpu")
     parser.add_argument("--data_dir", default=DATA_DIR)