sdil: add feedback-first apical calibration

1 files changed, 74 insertions, 5 deletions

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")