fix: model somato-dendritic baseline per neuron

1 files changed, 37 insertions, 3 deletions

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:")