sdil: add feedback-first apical calibration

1 files changed, 54 insertions, 11 deletions

diff --git a/sdil/core.py b/sdil/core.py
index 6536c17..20e94eb 100644
--- a/sdil/core.py
+++ b/sdil/core.py
@@ -482,6 +482,59 @@ class SDILConfig:
         self.direct_node_pert = direct_node_pert
 
 
+def _update_apical_vectorizer(net, h, c, r_list, qs, eta_A):
+    """Apply the local causal-regression update to all apical pathways."""
+    B = c.shape[0]
+    context = h[-2]
+    for l in range(net.L - 1):
+        calibration_error = qs[l] - r_list[l]
+        dA = calibration_error.t() @ c / B
+        net.A[l] += eta_A * dA
+        if net.vectorizer_mode == "soma_gated":
+            dgate = (calibration_error * torch.tanh(h[l + 1])).t() @ c / B
+            net.A_gate[l] += eta_A * dgate
+        elif net.vectorizer_mode == "context_gated":
+            features = (c.unsqueeze(2) * torch.tanh(context).unsqueeze(1)).flatten(1)
+            dgate = calibration_error.t() @ features / B
+            net.A_gate[l] += eta_A * dgate
+
+
+def apical_calibration_step(net, x, y, y_onehot, cfg,
+                            pert_mode=None, pert_ndirs=None):
+    """Calibrate the apical vectorizer while keeping all forward weights fixed.
+
+    This is the feedback-first half of a two-timescale protocol. It uses labels
+    and forward-only causal interventions exactly like an ordinary perturbation
+    event, but performs no hidden or output weight update. The caller controls
+    and accounts for the number of prefix steps and perturbation work.
+    """
+    if not cfg.learn_A:
+        raise ValueError("apical calibration requires learn_A=True")
+    if cfg.feedback != "error":
+        raise ValueError("feedback-first calibration currently requires error feedback")
+    mode = cfg.pert_mode if pert_mode is None else pert_mode
+    directions = cfg.pert_ndirs if pert_ndirs is None else pert_ndirs
+    if mode not in ("layerwise", "simultaneous"):
+        raise ValueError(f"unknown perturbation mode: {mode}")
+    if directions < 1:
+        raise ValueError("perturbation directions must be positive")
+
+    with torch.no_grad():
+        fwd = net.forward(x)
+        h = fwd["h"]
+        logits = h[-1]
+        loss = loss_ce(logits, y).mean().item()
+        c = net.output_error(logits, y_onehot)
+        context = h[-2]
+        r_list = [teaching_signal(net, l, c, h[l + 1], cfg, context)[0]
+                  for l in range(net.L - 1)]
+        estimator = (simultaneous_node_perturbation_targets
+                     if mode == "simultaneous" else node_perturbation_targets)
+        qs = estimator(net, x, y, sigma=cfg.pert_sigma, n_dirs=directions)
+        _update_apical_vectorizer(net, h, c, r_list, qs, cfg.eta_A)
+    return loss, {"error": c, "did_pert": True}
+
+
 def sdil_step(net, x, y, y_onehot, cfg, step, prev_error=None):
     """One SDIL training step (minibatch). Returns (loss, aux dict). No autograd."""
     with torch.no_grad():
@@ -552,17 +605,7 @@ def sdil_step(net, x, y, y_onehot, cfg, step, prev_error=None):
 
         # ================= apical vectorizer A via node perturbation =======
         if did_pert and cfg.learn_A:
-            for l in range(net.L - 1):
-                calibration_error = qs[l] - r_list[l]
-                dA = calibration_error.t() @ c / B     # (n_l, n_classes)
-                net.A[l] += cfg.eta_A * dA
-                if net.vectorizer_mode == "soma_gated":
-                    dgate = (calibration_error * torch.tanh(h[l + 1])).t() @ c / B
-                    net.A_gate[l] += cfg.eta_A * dgate
-                elif net.vectorizer_mode == "context_gated":
-                    features = (c.unsqueeze(2) * torch.tanh(context).unsqueeze(1)).flatten(1)
-                    dgate = calibration_error.t() @ features / B
-                    net.A_gate[l] += cfg.eta_A * dgate
+            _update_apical_vectorizer(net, h, c, r_list, qs, cfg.eta_A)
 
         # ================= predictor P (neutral) ==========================
         # KEY identification condition. P must learn the soma->apical coupling