From 5255166d13c46402deb2b43cfeafd8953c787fec Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Tue, 21 Jul 2026 06:46:33 -0500 Subject: fix: model somato-dendritic baseline per neuron --- experiments/run.py | 6 ++++-- experiments/smoke.py | 40 +++++++++++++++++++++++++++++++--- sdil/core.py | 61 +++++++++++++++++++++++++++++++++++++++------------- sdil/probes.py | 2 ++ 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] -- cgit v1.2.3