feat: model endogenous mixed apical traffic

1 files changed, 64 insertions, 15 deletions

diff --git a/sdil/core.py b/sdil/core.py
index d7be8c3..37c7165 100644
--- a/sdil/core.py
+++ b/sdil/core.py
@@ -100,7 +100,7 @@ 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"):
+                 predictor_mode="diagonal", traffic_mode="soma"):
         # sizes = [n_in, n_h1, ..., n_hk, n_out]
         self.sizes = list(sizes)
         self.L = len(sizes) - 1              # number of weight layers
@@ -109,6 +109,9 @@ class SDILNet:
         self.dtype = dtype
         self.act, self.act_prime = ACTS[act]
         self.feedback = feedback
+        if traffic_mode not in ("none", "soma", "topdown", "mixed"):
+            raise ValueError(f"unknown traffic_mode: {traffic_mode}")
+        self.traffic_mode = traffic_mode
         # residual (skip) connections on the constant-width hidden blocks: the
         # modern way to train deep nets with NO normalization.
         # h_{l+1}=h_l+alpha*phi(u_l) for the interior blocks (input projection and
@@ -140,10 +143,18 @@ class SDILNet:
         # population map.  "full" preserves the original experimental ablation.
         self.nuis_rho = nuis_rho
         gn = torch.Generator(device="cpu").manual_seed(nuis_seed)
+        gc = torch.Generator(device="cpu").manual_seed(nuis_seed + 1)
         if predictor_mode == "diagonal":
             # Positive log-normal slopes model ordinary soma-dendrite coupling.
             self.Bnuis = [torch.exp(0.25 * torch.randn(sizes[i + 1], generator=gn)).to(
                 device=device, dtype=dtype) for i in range(self.L - 1)]
+            # A second, independently generated traffic family carries the
+            # network's own high-level state back along corresponding feature
+            # channels. In constant-width residual networks this models normal
+            # contextual/top-down traffic rather than a target constructed from
+            # the local predictor. P still sees only the same cell's soma.
+            self.Bcontext = [torch.exp(0.25 * torch.randn(sizes[i + 1], generator=gc)).to(
+                device=device, dtype=dtype) for i in range(self.L - 1)]
             self.P = [torch.zeros(sizes[i + 1], device=device, dtype=dtype)
                       for i in range(self.L - 1)]
             self.P_bias = [torch.zeros(sizes[i + 1], device=device, dtype=dtype)
@@ -151,6 +162,9 @@ class SDILNet:
         else:
             self.Bnuis = [he((sizes[i + 1], sizes[i + 1]), 1.0, gen=gn)
                           for i in range(self.L - 1)]
+            context_dim = sizes[-2]
+            self.Bcontext = [he((sizes[i + 1], context_dim), 1.0, gen=gc)
+                             for i in range(self.L - 1)]
             self.P = [torch.zeros((sizes[i + 1], sizes[i + 1]), device=device, dtype=dtype)
                       for i in range(self.L - 1)]
             self.P_bias = None
@@ -195,14 +209,45 @@ class SDILNet:
         p = torch.softmax(logits, dim=1)
         return p - y_onehot
 
-    def apical(self, l, c, h_l=None):
-        """a_l = A_l c (+ soma-predictable nuisance rho * Bnuis h_l)."""
-        a = c @ self.A[l].t()
-        if self.nuis_rho and h_l is not None:
+    def apical_traffic(self, l, h_l, context=None):
+        """Return ordinary non-teaching traffic in the apical compartment.
+
+        ``soma`` models back-propagating somatic activity / normal
+        soma-dendrite coupling. ``topdown`` uses the network's own final hidden
+        state, generated independently of P, as contextual feedback. ``mixed``
+        combines both at matched total scale. The legacy ``nuis_rho`` argument
+        is retained as the traffic amplitude so existing result protocols are
+        exactly backward compatible.
+        """
+        if not self.nuis_rho or self.traffic_mode == "none":
+            return torch.zeros_like(h_l)
+
+        if self.predictor_mode == "diagonal":
+            soma = self.Bnuis[l] * h_l
+        else:
+            soma = h_l @ self.Bnuis[l].t()
+
+        if self.traffic_mode == "soma":
+            traffic = soma
+        else:
+            if context is None:
+                raise ValueError("topdown apical traffic requires a context state")
             if self.predictor_mode == "diagonal":
-                a = a + self.nuis_rho * self.Bnuis[l] * h_l
+                if context.shape[1] != h_l.shape[1]:
+                    raise ValueError(
+                        "diagonal topdown traffic requires constant hidden width; "
+                        f"got context={context.shape[1]} and layer={h_l.shape[1]}")
+                topdown = self.Bcontext[l] * context
             else:
-                a = a + self.nuis_rho * (h_l @ self.Bnuis[l].t())
+                topdown = context @ self.Bcontext[l].t()
+            traffic = topdown if self.traffic_mode == "topdown" else (soma + topdown) / _SQRT2
+        return self.nuis_rho * traffic
+
+    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()
+        if h_l is not None:
+            a = a + self.apical_traffic(l, h_l, context)
         return a
 
     def baseline(self, l, h_l):
@@ -211,15 +256,15 @@ class SDILNet:
             return self.P[l] * h_l + self.P_bias[l]
         return h_l @ self.P[l].t()
 
-    def innovation(self, l, c, h_l, use_residual=True):
+    def innovation(self, l, c, h_l, use_residual=True, context=None):
         """r_l = a_l - ahat_l  (or raw a_l if use_residual=False -> ablation)."""
-        a = self.apical(l, c, h_l)
+        a = self.apical(l, c, h_l, context)
         if use_residual:
             return a - self.baseline(l, h_l), a
         return a, a
 
 
-def teaching_signal(net, l, c, h_l, cfg):
+def teaching_signal(net, l, c, h_l, cfg, context=None):
     """Select the signal that drives a hidden layer's local plasticity.
 
     ``match_innovation_norm`` is a deliberately strong raw-apical control: it
@@ -231,7 +276,7 @@ def teaching_signal(net, l, c, h_l, cfg):
     Returns ``(teaching, raw_apical, innovation)`` so training and diagnostic
     probes use one definition of every signal.
     """
-    a = net.apical(l, c, h_l)
+    a = net.apical(l, c, h_l, context)
     innovation = a - net.baseline(l, h_l)
     if cfg.use_residual:
         teaching = innovation
@@ -416,10 +461,12 @@ def sdil_step(net, x, y, y_onehot, cfg, step, prev_error=None):
         if cfg.settle_steps > 0:
             h = _settle(net, h, u, c, cfg)
 
+        context = h[-2]
+
         # ---- compute innovations for hidden layers ----
         r_list, a_list = [], []
         for l in range(net.L - 1):
-            r, a, _ = teaching_signal(net, l, c, h[l + 1], cfg)
+            r, a, _ = teaching_signal(net, l, c, h[l + 1], cfg, context)
             r_list.append(r)
             a_list.append(a)
 
@@ -476,7 +523,7 @@ def sdil_step(net, x, y, y_onehot, cfg, step, prev_error=None):
             zero_c = torch.zeros_like(c)
             for l in range(net.L - 1):
                 if cfg.p_update_on_neutral:
-                    target = net.apical(l, zero_c, h[l + 1])   # nuisance only
+                    target = net.apical(l, zero_c, h[l + 1], context)  # traffic only
                 else:
                     target = a_list[l]                         # full apical (eats signal)
                 resid = target - net.baseline(l, h[l + 1])
@@ -494,11 +541,12 @@ def neutral_p_update(net, x, eta):
     with torch.no_grad():
         fwd = net.forward(x)
         h = fwd["h"]
+        context = h[-2]
         B = x.shape[0]
         zc = torch.zeros(B, net.n_classes, device=x.device, dtype=x.dtype)
         for l in range(net.L - 1):
             hl = h[l + 1]
-            target = net.apical(l, zc, hl)         # nuisance only (c=0)
+            target = net.apical(l, zc, hl, context)  # ordinary traffic only (c=0)
             resid = target - net.baseline(l, hl)
             _update_predictor(net, l, resid, hl, eta)
 
@@ -524,8 +572,9 @@ def _settle(net, h, u, c, cfg):
     h = list(h)
     dt_over_tau = 1.0 / max(1, cfg.settle_steps)
     for _ in range(cfg.settle_steps):
+        context = h[-2]
         for l in range(net.L - 1):
-            r, _, _ = teaching_signal(net, l, c, h[l + 1], cfg)
+            r, _, _ = teaching_signal(net, l, c, h[l + 1], cfg, context)
             target = net.act(u[l]) + cfg.kappa * r
             h[l + 1] = h[l + 1] + dt_over_tau * (-h[l + 1] + target)
         # recompute downstream pre-activations from settled hidden states