sdil: add feedback-first apical calibration

3 files changed, 107 insertions, 6 deletions

diff --git a/experiments/protocol_smoke.py b/experiments/protocol_smoke.py
index 18de4c9..7ee9b4b 100644
--- a/experiments/protocol_smoke.py
+++ b/experiments/protocol_smoke.py
@@ -59,6 +59,9 @@ def main():
         net, SDILConfig(pert_ndirs=2, pert_mode="simultaneous"))
     layerwise = calibration_work_per_event(
         net, SDILConfig(pert_ndirs=2, pert_mode="layerwise"))
+    overridden = calibration_work_per_event(
+        net, SDILConfig(pert_ndirs=1, pert_mode="simultaneous"),
+        pert_mode="layerwise", pert_ndirs=2)
     assert simultaneous == {
         "batch_loss_evaluations": 4,
         "forward_equivalent_batches": 5.0,
@@ -66,6 +69,7 @@ def main():
     }
     assert layerwise["batch_loss_evaluations"] == 9
     assert abs(layerwise["forward_equivalent_batches"] - 11.0 / 3.0) < 1e-12
+    assert overridden == layerwise
 
     lesion_net = SDILNet([1, 4, 4, 4, 4, 2], act="relu", residual=True, device="cpu")
     lesion_x = torch.linspace(0, 1, 16).view(-1, 1)
diff --git a/experiments/run.py b/experiments/run.py
index ea19366..ced7585 100644
--- a/experiments/run.py
+++ b/experiments/run.py
@@ -22,7 +22,8 @@ import torch
 import torch.nn.functional as F
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-from sdil.core import SDILNet, SDILConfig, sdil_step, neutral_p_update
+from sdil.core import (SDILNet, SDILConfig, apical_calibration_step, sdil_step,
+                       neutral_p_update)
 from sdil.baselines import BPNet, dfa_config, evaluate
 from sdil.local_baselines import FANet
 from sdil import probes
@@ -156,15 +157,16 @@ def residual_lesion_report(net, evaluation_loader, probe_x, fraction):
     }
 
 
-def calibration_work_per_event(net, cfg):
+def calibration_work_per_event(net, cfg, pert_mode=None, pert_ndirs=None):
     """Hardware-independent causal-calibration work for one minibatch.
 
     Forward equivalents use affine multiply-add work as the denominator. The
     returned loss evaluations count calls producing per-example scalar losses;
     the caller multiplies by the actual minibatch size for scalar observations.
     """
-    directions = cfg.pert_ndirs
-    if cfg.pert_mode == "simultaneous":
+    directions = cfg.pert_ndirs if pert_ndirs is None else pert_ndirs
+    mode = cfg.pert_mode if pert_mode is None else pert_mode
+    if mode == "simultaneous":
         return {
             "batch_loss_evaluations": 2 * directions,
             # The implementation performs one duplicate clean forward plus one
@@ -341,6 +343,8 @@ def train(args):
     calibration_example_loss_evaluations = 0
     calibration_forward_equivalent_examples = 0.0
     perturbation_batch_expansion = 0
+    feedback_warmup_examples = 0
+    feedback_warmup_events = 0
 
     # predictor warmup on neutral-period (c=0) drive, so P cancels the apical
     # nuisance before task plasticity relies on the residual (no-op when rho=0).
@@ -361,6 +365,63 @@ def train(args):
         warmup_wall_s = time.time() - warmup_t0
     else:
         warmup_wall_s = 0.0
+
+    # Feedback-first timescale separation. Causal perturbations fit A on a
+    # stationary forward network before noisy predictions can move W into a
+    # bad basin. The output readout is frozen as well. This is supervised
+    # calibration work, not free initialization, so every forward and scalar
+    # loss observation is included in the same hardware-independent ledger.
+    if args.a_warmup_steps > 0:
+        if args.mode != "sdil" or not args.learn_A:
+            raise ValueError("--a_warmup_steps requires SDIL with --learn_A 1")
+        device_sync(device)
+        feedback_warmup_t0 = time.time()
+        loader_state = (train_loader.g.get_state().clone()
+                        if hasattr(train_loader, "g") else None)
+        rng_devices = ([torch.cuda.current_device()]
+                       if str(device).startswith("cuda") and torch.cuda.is_available()
+                       else [])
+        # Warmup length must not silently change the minibatch order or random
+        # directions used by the subsequent joint phase. fork_rng and restoring
+        # the loader generator isolate those nuisance differences while keeping
+        # the learned A parameters.
+        try:
+            with torch.random.fork_rng(devices=rng_devices):
+                it = iter(train_loader)
+                for _ in range(args.a_warmup_steps):
+                    try:
+                        wx, wy = next(it)
+                    except StopIteration:
+                        it = iter(train_loader)
+                        wx, wy = next(it)
+                    wx, wy = wx.to(device), wy.to(device)
+                    wyoh = onehot(wy, n_out, device=device)
+                    apical_calibration_step(
+                        net, wx, wy, wyoh, cfg,
+                        pert_mode=args.a_warmup_mode,
+                        pert_ndirs=args.a_warmup_ndirs)
+                    feedback_warmup_examples += wx.shape[0]
+                    feedback_warmup_events += 1
+                    event = calibration_work_per_event(
+                        net, cfg, pert_mode=args.a_warmup_mode,
+                        pert_ndirs=args.a_warmup_ndirs)
+                    perturbation_events += 1
+                    calibration_batch_loss_evaluations += event["batch_loss_evaluations"]
+                    calibration_example_loss_evaluations += (
+                        event["batch_loss_evaluations"] * wx.shape[0])
+                    calibration_forward_equivalent_examples += (
+                        event["forward_equivalent_batches"] * wx.shape[0])
+                    perturbation_batch_expansion = max(
+                        perturbation_batch_expansion,
+                        event["perturbation_batch_expansion"])
+        finally:
+            if loader_state is not None:
+                train_loader.g.set_state(loader_state)
+        device_sync(device)
+        feedback_warmup_wall_s = time.time() - feedback_warmup_t0
+    else:
+        feedback_warmup_wall_s = 0.0
+
     for epoch in range(args.epochs):
         device_sync(device)
         train_t0 = time.time()
@@ -524,6 +585,7 @@ def train(args):
     log["final"]["wall_s"] = time.time() - t0
     log["timing"] = {
         "warmup_wall_s": warmup_wall_s,
+        "feedback_warmup_wall_s": feedback_warmup_wall_s,
         "training_loop_wall_s": train_wall_s,
         "diagnostics_wall_s": diagnostics_wall_s,
         "inline_diagnostics_wall_s": inline_diagnostics_wall_s,
@@ -535,12 +597,14 @@ def train(args):
         "train_steps": step,
         "ordinary_training_forward_examples": ordinary_forward_examples,
         "predictor_warmup_forward_examples": warmup_examples,
+        "feedback_warmup_forward_examples": feedback_warmup_examples,
+        "feedback_warmup_perturbation_events": feedback_warmup_events,
         "perturbation_events": perturbation_events,
         "calibration_batch_loss_evaluations": calibration_batch_loss_evaluations,
         "calibration_example_loss_evaluations": calibration_example_loss_evaluations,
         "calibration_forward_equivalent_examples": calibration_forward_equivalent_examples,
         "training_forward_equivalent_examples": (
-            ordinary_forward_examples + warmup_examples
+            ordinary_forward_examples + warmup_examples + feedback_warmup_examples
             + calibration_forward_equivalent_examples),
         "max_perturbation_batch_expansion": perturbation_batch_expansion,
     }
@@ -602,6 +666,11 @@ def get_args():
     p.add_argument("--p_neutral", type=int, default=1)      # P update on neutral (c=0) drive
     p.add_argument("--p_warmup_steps", type=int, default=200)  # pre-task neutral P warmup
     p.add_argument("--p_warmup_eta", type=float, default=0.05)
+    p.add_argument("--a_warmup_steps", type=int, default=0,
+                   help="A-only causal-calibration prefix with all forward weights frozen")
+    p.add_argument("--a_warmup_ndirs", type=int, default=4)
+    p.add_argument("--a_warmup_mode", default="layerwise",
+                   choices=["layerwise", "simultaneous"])
     p.add_argument("--nuis_rho", type=float, default=0.0)
     p.add_argument("--traffic_seed", type=int, default=1234,
                    help="ordinary apical-traffic projection seed; independent of model/data seeds")
diff --git a/experiments/smoke.py b/experiments/smoke.py
index a49e4c6..9971800 100644
--- a/experiments/smoke.py
+++ b/experiments/smoke.py
@@ -6,7 +6,7 @@ import sys
 import torch
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-from sdil.core import (SDILNet, SDILConfig, sdil_step,
+from sdil.core import (SDILNet, SDILConfig, apical_calibration_step, sdil_step,
                        node_perturbation_targets, simultaneous_node_perturbation_targets,
                        neutral_p_update, teaching_signal)
 from sdil.baselines import dfa_config
@@ -188,6 +188,33 @@ def check_direct_node_perturbation():
     print("CHECK0f direct node perturbation: exact q update; A untouched")
 
 
+def check_feedback_first_calibration():
+    """A-only calibration must leave the entire forward network unchanged."""
+    torch.manual_seed(23)
+    x = torch.randn(32, 5)
+    y = torch.randint(0, 3, (32,))
+    yoh = onehot(y, 3)
+    net = SDILNet([5, 7, 7, 3], act="relu", device="cpu", seed=11,
+                  residual=True, vectorizer_mode="context_gated")
+    weights_before = [weight.clone() for weight in net.W]
+    biases_before = [bias.clone() for bias in net.b]
+    apical_before = [weight.clone() for weight in net.A]
+    gates_before = [weight.clone() for weight in net.A_gate]
+    cfg = SDILConfig(eta_A=0.02, learn_A=True, learn_P=True,
+                     pert_ndirs=1, pert_mode="simultaneous")
+    apical_calibration_step(
+        net, x, y, yoh, cfg, pert_mode="simultaneous", pert_ndirs=2)
+    assert all(torch.equal(before, after)
+               for before, after in zip(weights_before, net.W))
+    assert all(torch.equal(before, after)
+               for before, after in zip(biases_before, net.b))
+    assert any(not torch.equal(before, after)
+               for before, after in zip(apical_before, net.A))
+    assert any(not torch.equal(before, after)
+               for before, after in zip(gates_before, net.A_gate))
+    print("CHECK0g feedback-first calibration: A changed; W/readout frozen")
+
+
 def main():
     torch.manual_seed(0)
     dev = "cpu"
@@ -197,6 +224,7 @@ def main():
     check_traffic_seed_isolation()
     check_state_conditioned_vectorizers()
     check_direct_node_perturbation()
+    check_feedback_first_calibration()
     print("loading MNIST subset...")
     tr, te, n_in, n_out = get_dataset("mnist", batch_size=128, device=dev)
     xb, yb = next(iter(tr))