""" SDIL -- Somato-Dendritic Innovation Learning. A non-backprop learning rule inspired by Harnett et al. 2026 (Nature), "Vectorized instructive signals in cortical dendrites". The load-bearing idea, faithful to the paper, is NOT "apical dendrites carry error". That is old (Guerguiev/Sacramento/Payeur/burst credit assignment). The new algorithmic object is the *somato-dendritic residual*: the teaching signal is the apical feedback MINUS the part of it that is predictable from the cell's own somatic activity. Only the innovation drives plasticity. For hidden layer l: u_l = W_l h_{l-1}, h_l = phi(u_l) (somatic / basal forward) a_l = A_l c_l (apical feedback, c_l = top-down signal) ahat_l = P_l h_l (predictable baseline: soma -> dendrite regression) r_l = a_l - ahat_l (INNOVATION = teaching signal) Three learning rules, all local (no weight transport, no reverse-mode graph): dW_l = eta * (r_l ⊙ phi'(u_l)) h_{l-1}^T (three-factor: pre * gain * innovation) dP_l = eta_P * (a_l - P_l h_l) h_l^T (slow; only on neutral periods) eta_P << eta dA_l = eta_A * (q_l - r_l) c_l^T (on perturbation trials; q_l = causal node-pert estimate) q_l is an amortized node-perturbation estimate of the causal descent direction -grad_{h_l} L, obtained by occasionally perturbing h_l and watching the loss. Learning A_l this way (rather than fixing it random) is what keeps SDIL from degenerating into DFA. This module deliberately does NOT use autograd for learning. Autograd is used ONLY inside probes.py to MEASURE the true gradient for alignment diagnostics. """ import math import torch # -------------------------------------------------------------------------- # activations # -------------------------------------------------------------------------- def tanh(u): return torch.tanh(u) def tanh_prime(u): t = torch.tanh(u) return 1.0 - t * t _SQRT2 = 2.0 ** 0.5 _INV_SQRT_2PI = 0.3989422804014327 def gelu(u): return 0.5 * u * (1.0 + torch.erf(u / _SQRT2)) def gelu_prime(u): return 0.5 * (1.0 + torch.erf(u / _SQRT2)) + u * _INV_SQRT_2PI * torch.exp(-0.5 * u * u) def silu(u): s = torch.sigmoid(u) return u * s def silu_prime(u): s = torch.sigmoid(u) return s * (1.0 + u * (1.0 - s)) def relu(u): return torch.relu(u) def relu_prime(u): return (u > 0).to(u.dtype) ACTS = {"tanh": (tanh, tanh_prime), "gelu": (gelu, gelu_prime), "silu": (silu, silu_prime), "relu": (relu, relu_prime)} # -------------------------------------------------------------------------- # model # -------------------------------------------------------------------------- class SDILNet: """Fully-connected SDIL network. Weights/parameters are plain tensors (requires_grad=False). Learning is manual and local. The output layer is trained with the exact local delta rule (softmax-CE gradient wrt logits IS the local output error), which is biologically unproblematic; only hidden layers use the SDIL innovation. feedback signal c: by default the output error e = softmax(logits) - onehot, broadcast to every hidden layer (dim = n_classes). With A_l learned this is "amortized node perturbation through apical dendrites"; with A_l fixed random it reduces to DFA (used as an ablation). """ 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, 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 self.n_classes = sizes[-1] self.device = device self.dtype = dtype self.act, self.act_prime = ACTS[act] self.feedback = feedback # residual (skip) connections on the constant-width hidden blocks: the # modern way to train deep nets with NO normalization. # h_{l+1}=h_l+alpha*phi(u_l) for the interior blocks (input projection and # readout stay plain). alpha=1/sqrt(#hidden) (SkipInit/Fixup/LayerScale) # keeps activation variance bounded across depth without any BN. self.residual = residual n_hidden = self.L - 1 self.res_alpha = (1.0 / math.sqrt(max(1, n_hidden))) if residual else 1.0 g = torch.Generator(device="cpu").manual_seed(seed) def he(shape, scale, gen=g): fan_in = shape[1] return (torch.randn(*shape, generator=gen) * (scale / math.sqrt(fan_in))).to(device=device, dtype=dtype) # forward weights + bias self.W = [he((sizes[i + 1], sizes[i]), w_scale) for i in range(self.L)] self.b = [torch.zeros(sizes[i + 1], device=device, dtype=dtype) for i in range(self.L)] # 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)] 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 ---- # 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) 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] self.mb = [torch.zeros_like(bb) for bb in self.b] # ------- forward ------- def forward(self, x, settle_steps=0, kappa=0.0, c_override=None): """Return dict of activations. h[0]=x; for l in 1..L-1 hidden tanh; logits = W_L h_{L-1} + b_L (linear). If settle_steps>0, run online apical control: after the feedforward pass, iterate h_l <- h_l + (1/tau)(-h_l + phi(u_l) + kappa r_l) using the current innovation r_l (desired-velocity interpretation). Requires labels via c_override or is applied post-hoc by the caller. """ h = [x] u = [] for l in range(self.L): pre = h[-1] ul = pre @ self.W[l].t() + self.b[l] u.append(ul) if l < self.L - 1: act = self.act(ul) # skip on interior blocks only (l>=1 => width->width, dims match) if self.residual and l >= 1: h.append(pre + self.res_alpha * act) else: h.append(act) else: h.append(ul) # logits (linear output) return {"h": h, "u": u} def logits(self, x): return self.forward(x)["h"][-1] # ------- feedback signal ------- def output_error(self, logits, y_onehot): # softmax-CE gradient wrt logits p = torch.softmax(logits, dim=1) return p - y_onehot def apical(self, l, c, h_l=None): """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: 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): """r_l = a_l - ahat_l (or raw a_l if use_residual=False -> ablation).""" a = self.apical(l, c, h_l) if use_residual: return a - self.baseline(l, h_l), a return a, a # -------------------------------------------------------------------------- # node perturbation: causal estimate of -grad_{h_l} L used to TRAIN A_l # -------------------------------------------------------------------------- def loss_ce(logits, y): return torch.nn.functional.cross_entropy(logits, y, reduction="none") def node_perturbation_targets(net, x, y, sigma=1e-2, antithetic=True, n_dirs=1): """Estimate q_l ~ -grad_{h_l}L for every hidden layer by forward-only perturbation of the somatic activity. Returns list over hidden layers of (B, n_l) tensors. No autograd. We perturb h_l -> h_l + sigma*xi, propagate forward THROUGH THE REST of the net (weights only), read the loss change dL, and form the antithetic estimate q_l = -((L(+sigma*xi) - L(-sigma*xi)) / (2*sigma)) * xi Antithetic pairs (+xi, -xi) cancel the O(sigma^2) even term. Averaging over n_dirs independent directions cuts the estimator variance ~1/n_dirs, which the predecessor project's failures flagged as the main risk of learning a feedback pathway from causal probes. """ with torch.no_grad(): base = net.forward(x) h = base["h"] L0 = loss_ce(h[-1], y) # (B,) qs = [] for l in range(net.L - 1): # hidden layers hl = h[l + 1] acc = torch.zeros_like(hl) for _ in range(n_dirs): xi = torch.randn_like(hl) dL = _relayer_loss(net, h, l + 1, hl + sigma * xi, y) - L0 # Correctly-scaled zeroth-order estimator of grad_{h_l}L: with # perturbation delta=sigma*xi, ghat = delta*dL/sigma^2 = xi*dL/sigma, # so E[ghat]=grad. Dividing by sigma (NOT sigma^2) makes ||q||~||grad|| # -> the three-factor update is genuine SGD at learning rate eta. if antithetic: dL2 = _relayer_loss(net, h, l + 1, hl - sigma * xi, y) - L0 coeff = (dL - dL2) / (2.0 * sigma) else: coeff = dL / sigma acc += -(coeff.unsqueeze(1)) * xi qs.append(acc / n_dirs) # (B, n_l) return qs def _relayer_loss(net, h, start_idx, h_start_new, y): """Recompute loss when hidden activation at layer index `start_idx` (1-based into h) is replaced by h_start_new, propagating forward. start_idx>=1 so every layer in the tail is an interior/readout layer.""" cur = h_start_new for l in range(start_idx, net.L): pre = cur ul = pre @ net.W[l].t() + net.b[l] if l < net.L - 1: act = net.act(ul) cur = pre + net.res_alpha * act if net.residual else act # interior blocks skip (l>=1) else: cur = ul return loss_ce(cur, y) # -------------------------------------------------------------------------- # the SDIL update # -------------------------------------------------------------------------- class SDILConfig: def __init__(self, eta=0.05, eta_A=0.02, eta_P=0.002, use_residual=True, learn_A=True, learn_P=True, pert_sigma=1e-2, pert_every=5, pert_ndirs=1, momentum=0.0, wd=0.0, settle_steps=0, kappa=0.0, feedback="error", p_update_on_neutral=True, normalize_delta=False): self.eta = eta self.eta_A = eta_A self.eta_P = eta_P self.use_residual = use_residual # ablation: residual vs raw apical self.learn_A = learn_A # ablation: perturbation-trained vs fixed-random A self.learn_P = learn_P # ablation: predictor on/off self.pert_sigma = pert_sigma self.pert_every = pert_every # run node-perturbation every k steps self.pert_ndirs = pert_ndirs # directions averaged per perturbation self.momentum = momentum self.wd = wd self.settle_steps = settle_steps # online apical control iterations self.kappa = kappa self.feedback = feedback # "error" | "error_deriv" self.p_update_on_neutral = p_update_on_neutral # node perturbation estimates the descent DIRECTION well but its # magnitude is noisy, so A can learn a mis-scaled per-layer step. # Normalising each layer's delta to unit RMS decouples step size (set by # eta) from that noise -- like normalized SGD. self.normalize_delta = normalize_delta def sdil_step(net, x, y, y_onehot, cfg, step, prev_error=None): """One SDIL training step (minibatch). Returns (loss, aux dict). No autograd.""" with torch.no_grad(): B = x.shape[0] fwd = net.forward(x) h, u = fwd["h"], fwd["u"] logits = h[-1] loss = loss_ce(logits, y).mean().item() e = net.output_error(logits, y_onehot) # (B, n_classes) # feedback signal c if cfg.feedback == "error_deriv" and prev_error is not None: c = e - prev_error # temporal error difference else: c = e # ---- optional online apical control (desired-velocity dynamics) ---- if cfg.settle_steps > 0: h = _settle(net, h, u, c, cfg) # ---- compute innovations for hidden layers ---- r_list, a_list = [], [] for l in range(net.L - 1): r, a = net.innovation(l, c, h[l + 1], use_residual=cfg.use_residual) r_list.append(r) a_list.append(a) # Calibrate A against the same network state that produced r and c. # Measuring q after mutating W would pair a post-update causal target # with a pre-update prediction, introducing an avoidable stale-target # error (especially at large learning rates). did_pert = cfg.learn_A and (step % cfg.pert_every == 0) qs = None if did_pert: qs = node_perturbation_targets( net, x, y, sigma=cfg.pert_sigma, n_dirs=cfg.pert_ndirs) # ================= forward weight updates ================= # hidden layers: three-factor local rule for l in range(net.L - 1): gain = net.act_prime(u[l]) # phi'(u_l) (B, n_l) delta = r_list[l] * gain # (B, n_l) if cfg.normalize_delta: delta = delta / (delta.pow(2).mean().sqrt() + 1e-8) # For an interior residual block h'=h+alpha*phi(Wh), the local # Jacobian dh'/dW contains alpha. Omitting it preserves direction # but makes the effective step grow as 1/alpha=sqrt(depth), which # confounds depth-scaling comparisons. The input projection (l=0) # is plain rather than residual and therefore has scale 1. if net.residual and l >= 1: delta = net.res_alpha * delta dW = delta.t() @ h[l] / B # (n_l, n_{l-1}) db = delta.mean(0) _apply(net, l, dW, db, cfg) # output layer: exact local delta rule (output error is locally available) gL = e # grad wrt logits dWL = -(gL.t() @ h[net.L - 1]) / B # descent dbL = -gL.mean(0) _apply(net, net.L - 1, dWL, dbL, cfg) # ================= apical vectorizer A via node perturbation ======= if did_pert: for l in range(net.L - 1): dA = (qs[l] - r_list[l]).t() @ c / B # (n_l, n_classes) net.A[l] += cfg.eta_A * dA # ================= predictor P (slow) ============================= # KEY timescale-separation condition. P must learn the soma->apical # coupling WITHOUT absorbing the teaching signal. On "neutral periods" # the apical compartment carries only the soma-predictable nuisance drive # (top-down teaching c is absent), so we fit P to the c=0 apical. This # makes P -> nuisance map (rho*Bnuis) and leaves the innovation intact. # With p_update_on_neutral=False we fit P to the full task-period apical # (ablation) -- that path demonstrably eats the teaching signal. if cfg.learn_P: zero_c = torch.zeros_like(c) for l in range(net.L - 1): if cfg.p_update_on_neutral: target = net.apical(l, zero_c, h[l + 1]) # nuisance only else: target = a_list[l] # full apical (eats signal) resid = target - net.baseline(l, h[l + 1]) _update_predictor(net, l, resid, h[l + 1], cfg.eta_P) return loss, {"error": e, "did_pert": did_pert} def neutral_p_update(net, x, eta): """Fit the predictor P to the neutral-period (c=0) apical drive, i.e. the soma-predictable nuisance. Used both as a pre-task warmup and (optionally) ongoing. Because the nuisance is linear in h, a converged P cancels it for any h, so warming P up before task plasticity prevents the nuisance from corrupting the forward weights early. No autograd.""" with torch.no_grad(): fwd = net.forward(x) h = fwd["h"] B = x.shape[0] zc = torch.zeros(B, net.n_classes, device=x.device, dtype=x.dtype) for l in range(net.L - 1): hl = h[l + 1] target = net.apical(l, zc, hl) # nuisance only (c=0) resid = target - net.baseline(l, hl) _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): """Online apical control: tau dh = -h + phi(u) + kappa r, a few Euler steps. Updates hidden states in place-ish (returns new h list).""" h = list(h) dt_over_tau = 1.0 / max(1, cfg.settle_steps) for _ in range(cfg.settle_steps): for l in range(net.L - 1): r, _ = net.innovation(l, c, h[l + 1], use_residual=cfg.use_residual) target = net.act(u[l]) + cfg.kappa * r h[l + 1] = h[l + 1] + dt_over_tau * (-h[l + 1] + target) # recompute downstream pre-activations from settled hidden states for l in range(1, net.L): u[l] = h[l] @ net.W[l].t() + net.b[l] if l < net.L - 1: pass # h[l+1] already being settled; keep u fresh for gains return h def _apply(net, l, dW, db, cfg): if cfg.wd: dW = dW - cfg.wd * net.W[l] if cfg.momentum: net.mW[l].mul_(cfg.momentum).add_(dW) net.mb[l].mul_(cfg.momentum).add_(db) net.W[l] += cfg.eta * net.mW[l] net.b[l] += cfg.eta * net.mb[l] else: net.W[l] += cfg.eta * dW net.b[l] += cfg.eta * db