fix: model somato-dendritic baseline per neuron

1 files changed, 46 insertions, 15 deletions

diff --git a/sdil/core.py b/sdil/core.py
index ce42650..048408e 100644
--- a/sdil/core.py
+++ b/sdil/core.py
@@ -99,7 +99,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):
+                 nuis_rho=0.0, nuis_seed=1234, residual=False,
+                 predictor_mode="diagonal"):
         # sizes = [n_in, n_h1, ..., n_hk, n_out]
         self.sizes = list(sizes)
         self.L = len(sizes) - 1              # number of weight layers
@@ -128,19 +129,31 @@ class SDILNet:
 
         # apical vectorizer A_l : c (n_classes) -> hidden velocity (n_l), hidden layers only (l=1..L-1)
         self.A = [he((sizes[i + 1], self.n_classes), a_scale) for i in range(self.L - 1)]
-        # predictor P_l : soma h_l (n_l) -> apical baseline (n_l)
-        self.P = [torch.zeros((sizes[i + 1], sizes[i + 1]), 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
 
         # ---- soma-predictable nuisance on the apical compartment ----
-        # Models the biological reality that apical dendrites also receive top-down
-        # drive correlated with the cell's own recent activity. This drive is
-        # PREDICTABLE from h_l (linear in h_l), so P_l can cancel it but A_l (which
-        # only sees the top-down signal c) cannot. rho controls contamination.
-        # It is the setting where the *residual* (innovation) matters: raw apical
-        # alignment degrades with rho, residual alignment should be preserved.
+        # Harnett et al. define each cell's SD residual using a least-squares line
+        # between that same cell's somatic and dendritic event magnitudes.  The
+        # faithful default is therefore a per-neuron affine predictor, not a dense
+        # population map.  "full" preserves the original experimental ablation.
         self.nuis_rho = nuis_rho
         gn = torch.Generator(device="cpu").manual_seed(nuis_seed)
-        self.Bnuis = [he((sizes[i + 1], sizes[i + 1]), 1.0, gen=gn) for i in range(self.L - 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)]
+            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)
+                           for i in range(self.L - 1)]
+        else:
+            self.Bnuis = [he((sizes[i + 1], sizes[i + 1]), 1.0, gen=gn)
+                          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
 
         # velocity buffers for momentum on forward weights (optional)
         self.mW = [torch.zeros_like(w) for w in self.W]
@@ -186,11 +199,16 @@ class SDILNet:
         """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:
-            a = a + self.nuis_rho * (h_l @ self.Bnuis[l].t())
+            if self.predictor_mode == "diagonal":
+                a = a + self.nuis_rho * self.Bnuis[l] * h_l
+            else:
+                a = a + self.nuis_rho * (h_l @ self.Bnuis[l].t())
         return a
 
     def baseline(self, l, h_l):
         """ahat_l = P_l h_l."""
+        if self.predictor_mode == "diagonal":
+            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):
@@ -368,8 +386,7 @@ def sdil_step(net, x, y, y_onehot, cfg, step, prev_error=None):
                 else:
                     target = a_list[l]                         # full apical (eats signal)
                 resid = target - net.baseline(l, h[l + 1])
-                dP = resid.t() @ h[l + 1] / B          # (n_l, n_l)
-                net.P[l] += cfg.eta_P * dP
+                _update_predictor(net, l, resid, h[l + 1], cfg.eta_P)
 
     return loss, {"error": e, "did_pert": did_pert}
 
@@ -389,8 +406,22 @@ def neutral_p_update(net, x, eta):
             hl = h[l + 1]
             target = net.apical(l, zc, hl)         # nuisance only (c=0)
             resid = target - net.baseline(l, hl)
-            dP = resid.t() @ hl / B
-            net.P[l] += eta * dP
+            _update_predictor(net, l, resid, hl, eta)
+
+
+def _update_predictor(net, l, resid, h_l, eta):
+    """Local least-squares update for the selected soma-dendrite predictor."""
+    if net.predictor_mode == "diagonal":
+        # Normalized LMS makes the per-cell slope learning rate independent of
+        # that cell's activity variance. Centering separates slope and intercept
+        # updates, matching an online least-squares line fit.
+        h_centered = h_l - h_l.mean(0)
+        resid_centered = resid - resid.mean(0)
+        variance = h_centered.pow(2).mean(0)
+        net.P[l] += eta * ((resid_centered * h_centered).mean(0) / (variance + 1e-6))
+        net.P_bias[l] += eta * resid.mean(0)
+    else:
+        net.P[l] += eta * (resid.t() @ h_l / h_l.shape[0])
 
 
 def _settle(net, h, u, c, cfg):