method: condition apical vectorizers on neural state

2 files changed, 52 insertions, 5 deletions

diff --git a/sdil/core.py b/sdil/core.py
index 8b3f923..e66dc39 100644
--- a/sdil/core.py
+++ b/sdil/core.py
@@ -102,7 +102,8 @@ class SDILNet:
     def __init__(self, sizes, act="tanh", device="cpu", seed=0,
                  w_scale=1.0, a_scale=1.0, feedback="error", dtype=torch.float32,
                  nuis_rho=0.0, nuis_seed=1234, residual=False,
-                 predictor_mode="diagonal", traffic_mode="soma"):
+                 predictor_mode="diagonal", traffic_mode="soma",
+                 vectorizer_mode="linear"):
         # sizes = [n_in, n_h1, ..., n_hk, n_out]
         self.sizes = list(sizes)
         self.L = len(sizes) - 1              # number of weight layers
@@ -132,8 +133,24 @@ class SDILNet:
         self.W = [he((sizes[i + 1], sizes[i]), w_scale) for i in range(self.L)]
         self.b = [torch.zeros(sizes[i + 1], device=device, dtype=dtype) for i in range(self.L)]
 
-        # apical vectorizer A_l : c (n_classes) -> hidden velocity (n_l), hidden layers only (l=1..L-1)
+        # Apical vectorizer A_l : c (n_classes) -> hidden velocity (n_l),
+        # hidden layers only.  A fixed linear map cannot represent credit that
+        # changes across activation regions. Optional zero-initialized gated
+        # terms preserve the initial linear model while letting perturbation
+        # calibration learn state-dependent feedback locally.
+        if vectorizer_mode not in ("linear", "soma_gated", "context_gated"):
+            raise ValueError(f"unknown vectorizer_mode: {vectorizer_mode}")
+        self.vectorizer_mode = vectorizer_mode
         self.A = [he((sizes[i + 1], self.n_classes), a_scale) for i in range(self.L - 1)]
+        if vectorizer_mode == "linear":
+            self.A_gate = None
+        elif vectorizer_mode == "soma_gated":
+            self.A_gate = [torch.zeros_like(weight) for weight in self.A]
+        else:
+            context_dim = sizes[-2]
+            self.A_gate = [torch.zeros((sizes[i + 1], self.n_classes * context_dim),
+                                       device=device, dtype=dtype)
+                           for i in range(self.L - 1)]
         if predictor_mode not in ("diagonal", "full"):
             raise ValueError(f"unknown predictor_mode: {predictor_mode}")
         self.predictor_mode = predictor_mode
@@ -245,9 +262,29 @@ class SDILNet:
             traffic = topdown if self.traffic_mode == "topdown" else (soma + topdown) / _SQRT2
         return self.nuis_rho * traffic
 
+    def vectorizer(self, l, c, h_l=None, context=None):
+        """Return the learned instructional component of apical activity.
+
+        ``soma_gated`` adds ``h_i G_i c`` independently in every cell.
+        ``context_gated`` adds a linear readout of ``c outer context``. Both
+        vanish when c=0, so neutral-period predictor identification remains
+        unchanged.
+        """
+        value = c @ self.A[l].t()
+        if self.vectorizer_mode == "linear":
+            return value
+        if self.vectorizer_mode == "soma_gated":
+            if h_l is None:
+                raise ValueError("soma-gated vectorizer requires somatic activity")
+            return value + h_l * (c @ self.A_gate[l].t())
+        if context is None:
+            raise ValueError("context-gated vectorizer requires a context state")
+        features = (c.unsqueeze(2) * context.unsqueeze(1)).flatten(1)
+        return value + features @ self.A_gate[l].t()
+
     def apical(self, l, c, h_l=None, context=None):
         """a_l = A_l c + ordinary non-teaching apical traffic."""
-        a = c @ self.A[l].t()
+        a = self.vectorizer(l, c, h_l, context)
         if h_l is not None:
             a = a + self.apical_traffic(l, h_l, context)
         return a
@@ -510,8 +547,16 @@ def sdil_step(net, x, y, y_onehot, cfg, step, prev_error=None):
         # ================= apical vectorizer A via node perturbation =======
         if did_pert:
             for l in range(net.L - 1):
-                dA = (qs[l] - r_list[l]).t() @ c / B   # (n_l, n_classes)
+                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 * h[l + 1]).t() @ c / B
+                    net.A_gate[l] += cfg.eta_A * dgate
+                elif net.vectorizer_mode == "context_gated":
+                    features = (c.unsqueeze(2) * context.unsqueeze(1)).flatten(1)
+                    dgate = calibration_error.t() @ features / B
+                    net.A_gate[l] += cfg.eta_A * dgate
 
         # ================= predictor P (neutral) ==========================
         # KEY identification condition. P must learn the soma->apical coupling
diff --git a/sdil/probes.py b/sdil/probes.py
index 1698d75..0075bca 100644
--- a/sdil/probes.py
+++ b/sdil/probes.py
@@ -74,7 +74,7 @@ def alignment_report(net, x, y, y_onehot, cfg):
         negg = -g
         teaching, a, innovation = core.teaching_signal(
             net, l, c, h[l + 1], cfg, context)
-        Ac = c @ net.A[l].t()                           # pure vectorizer output
+        Ac = net.vectorizer(l, c, h[l + 1], context)    # pure vectorizer output
         traffic = net.apical_traffic(l, h[l + 1], context)
         unexplained = traffic - net.baseline(l, h[l + 1])
         # include the phi' gain that the plasticity rule actually applies,
@@ -145,6 +145,8 @@ def _clone(net):
     n.W = [w.clone() for w in net.W]
     n.b = [b.clone() for b in net.b]
     n.A = [a.clone() for a in net.A]
+    n.A_gate = ([a.clone() for a in net.A_gate]
+                if getattr(net, "A_gate", None) is not None else None)
     n.P = [p.clone() for p in net.P]
     n.P_bias = ([p.clone() for p in net.P_bias]
                 if getattr(net, "P_bias", None) is not None else None)