fix: model somato-dendritic baseline per neuron

4 files changed, 89 insertions, 20 deletions

diff --git a/experiments/run.py b/experiments/run.py
index 3ad0bdd..bb2d151 100644
--- a/experiments/run.py
+++ b/experiments/run.py
@@ -45,13 +45,14 @@ def build(args, device):
     sizes = [args.n_in] + [args.width] * args.depth + [10]
     if args.mode == "bp":
         net = BPNet(sizes, act=args.act, device=device, seed=args.seed,
-                    w_scale=args.w_scale, nuis_rho=0.0, residual=bool(args.residual))
+                    w_scale=args.w_scale, nuis_rho=0.0, residual=bool(args.residual),
+                    predictor_mode=args.predictor_mode)
         cfg = SDILConfig(eta=args.eta, momentum=args.momentum)
         return net, cfg
     net = SDILNet(sizes, act=args.act, device=device, seed=args.seed,
                   w_scale=args.w_scale, a_scale=args.a_scale,
                   nuis_rho=args.nuis_rho, feedback=args.feedback,
-                  residual=bool(args.residual))
+                  residual=bool(args.residual), predictor_mode=args.predictor_mode)
     if args.mode == "dfa":
         cfg = dfa_config(eta=args.eta, momentum=args.momentum)
     elif args.mode == "sdil":
@@ -175,6 +176,7 @@ def get_args():
     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("--nuis_rho", type=float, default=0.0)
+    p.add_argument("--predictor_mode", default="diagonal", choices=["diagonal", "full"])
     p.add_argument("--normalize_delta", type=int, default=0)
     p.add_argument("--settle_steps", type=int, default=0)
     p.add_argument("--kappa", type=float, default=0.0)
diff --git a/experiments/smoke.py b/experiments/smoke.py
index 092e9f5..8196d4f 100644
--- a/experiments/smoke.py
+++ b/experiments/smoke.py
@@ -6,7 +6,8 @@ import sys
 import torch
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-from sdil.core import SDILNet, SDILConfig, sdil_step, node_perturbation_targets
+from sdil.core import (SDILNet, SDILConfig, sdil_step,
+                       node_perturbation_targets, neutral_p_update)
 from sdil.baselines import dfa_config
 from sdil import probes
 from sdil.data import get_dataset, onehot
@@ -53,10 +54,38 @@ def check_residual_local_jacobian():
         assert abs(norm_ratio - 1.0) < 1e-5
 
 
+def check_neutral_predictor():
+    """Per-neuron neutral regression should remove normal dendritic coupling."""
+    torch.manual_seed(11)
+    net = SDILNet([20, 32, 32, 5], device="cpu", seed=5, nuis_rho=2.0,
+                  predictor_mode="diagonal")
+    x = torch.randn(128, 20)
+    initial = []
+    final = []
+    with torch.no_grad():
+        h = net.forward(x)["h"]
+        zero_c = torch.zeros(x.shape[0], net.n_classes)
+        for l in range(net.L - 1):
+            target = net.apical(l, zero_c, h[l + 1])
+            initial.append((target - net.baseline(l, h[l + 1])).norm() / target.norm())
+    for _ in range(300):
+        neutral_p_update(net, x, 0.05)
+    with torch.no_grad():
+        h = net.forward(x)["h"]
+        for l in range(net.L - 1):
+            target = net.apical(l, zero_c, h[l + 1])
+            final.append((target - net.baseline(l, h[l + 1])).norm() / target.norm())
+    print("CHECK0b neutral predictor residual/target:")
+    for l, (before, after) in enumerate(zip(initial, final)):
+        print(f"  layer {l}: {before.item():.4f} -> {after.item():.4f}")
+        assert after < 0.03
+
+
 def main():
     torch.manual_seed(0)
     dev = "cpu"
     check_residual_local_jacobian()
+    check_neutral_predictor()
     print("loading MNIST subset...")
     tr, te, n_in, n_out = get_dataset("mnist", batch_size=128, device=dev)
     xb, yb = next(iter(tr))
@@ -113,12 +142,17 @@ def main():
     print("\nCHECK4 residualization under soma-predictable nuisance (rho=2.0):")
     net_n = SDILNet(sizes, device=dev, seed=4, nuis_rho=2.0)
     cfg_n = SDILConfig(eta=0.1, eta_A=0.05, eta_P=0.02, pert_every=2, pert_ndirs=8, momentum=0.9)
+    for _ in range(200):
+        neutral_p_update(net_n, x, 0.05)
     for step in range(401):
         sdil_step(net_n, x, y, yoh, cfg_n, step)
     al_n = probes.alignment_report(net_n, x, y, yoh, cfg_n)
-    print(f"  residual cos(r,-g)      = {sum(al_n['cos_r_negg'])/len(al_n['cos_r_negg']):+.3f}")
-    print(f"  raw apical cos(a,-g)    = {sum(al_n['cos_apical_negg'])/len(al_n['cos_apical_negg']):+.3f}")
+    residual_cos = sum(al_n['cos_r_negg']) / len(al_n['cos_r_negg'])
+    apical_cos = sum(al_n['cos_apical_negg']) / len(al_n['cos_apical_negg'])
+    print(f"  residual cos(r,-g)      = {residual_cos:+.3f}")
+    print(f"  raw apical cos(a,-g)    = {apical_cos:+.3f}")
     print(f"  pure A c   cos(Ac,-g)   = {sum(al_n['cos_Ac_negg'])/len(al_n['cos_Ac_negg']):+.3f}")
+    assert residual_cos > apical_cos + 0.1
 
     # ---- CHECK 5: single-step loss-decrease ratio vs exact GD ----
     print("\nCHECK5 single-step loss-decrease ratio vs exact GD:")
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):
diff --git a/sdil/probes.py b/sdil/probes.py
index 5489c6e..b58fe89 100644
--- a/sdil/probes.py
+++ b/sdil/probes.py
@@ -105,6 +105,8 @@ def _clone(net):
     n.b = [b.clone() for b in net.b]
     n.A = [a.clone() for a in net.A]
     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)
     n.Bnuis = [b.clone() for b in net.Bnuis]
     n.mW = [torch.zeros_like(w) for w in net.W]
     n.mb = [torch.zeros_like(b) for b in net.b]