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 --- sdil/core.py | 61 +++++++++++++++++++++++++++++++++++++++++++--------------- sdil/probes.py | 2 ++ 2 files changed, 48 insertions(+), 15 deletions(-) (limited to 'sdil') 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