""" SDIL -- Somato-Dendritic Innovation Learning. A non-backprop learning rule inspired by Harnett et al. 2026 (Nature), "Vectorized instructive signals in cortical dendrites". The Harnett-specific object studied here 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. Apical credit signals and learned feedback from node perturbations both have substantial prior art; see NOVELTY.md for the exact claim boundary. 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 (identified only on neutral periods) dA_l = eta_A * (q_l - r_l) c_l^T (on perturbation trials; q_l = causal node-pert estimate) q_l is a 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 follows the learned synthetic-feedback approach of Lansdell, Prakash & Kording (ICLR 2020). Relative to fixed DFA it supplies an adaptive feedback backbone; it is not itself claimed as an SDIL novelty. 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 P=0 and no traffic, learning A_l from node perturbations is the direct learned-feedback method of Lansdell et al.; with A_l fixed random it reduces to DFA. SDIL's candidate addition is the per-cell innovation under mixed apical traffic. """ 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", traffic_mode="soma", vectorizer_mode="linear"): # 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 if traffic_mode not in ("none", "soma", "topdown", "mixed"): raise ValueError(f"unknown traffic_mode: {traffic_mode}") self.traffic_mode = traffic_mode # 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. A fixed linear map cannot represent credit that # changes across activation regions. Optional zero-initialized gated # terms preserve the initial linear model while letting perturbation # calibration learn state-dependent feedback locally. if vectorizer_mode not in ("linear", "soma_gated", "context_gated"): raise ValueError(f"unknown vectorizer_mode: {vectorizer_mode}") self.vectorizer_mode = vectorizer_mode self.A = [he((sizes[i + 1], self.n_classes), a_scale) for i in range(self.L - 1)] if vectorizer_mode == "linear": self.A_gate = None elif vectorizer_mode == "soma_gated": self.A_gate = [torch.zeros_like(weight) for weight in self.A] else: context_dim = sizes[-2] self.A_gate = [torch.zeros((sizes[i + 1], self.n_classes * context_dim), 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 ---- # 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) gc = torch.Generator(device="cpu").manual_seed(nuis_seed + 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)] # A second, independently generated traffic family carries the # network's own high-level state back along corresponding feature # channels. In constant-width residual networks this models normal # contextual/top-down traffic rather than a target constructed from # the local predictor. P still sees only the same cell's soma. self.Bcontext = [torch.exp(0.25 * torch.randn(sizes[i + 1], generator=gc)).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)] context_dim = sizes[-2] self.Bcontext = [he((sizes[i + 1], context_dim), 1.0, gen=gc) 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_traffic(self, l, h_l, context=None): """Return ordinary non-teaching traffic in the apical compartment. ``soma`` models back-propagating somatic activity / normal soma-dendrite coupling. ``topdown`` uses the network's own final hidden state, generated independently of P, as contextual feedback. ``mixed`` combines both at matched total scale. The legacy ``nuis_rho`` argument is retained as the traffic amplitude so existing result protocols are exactly backward compatible. """ if not self.nuis_rho or self.traffic_mode == "none": return torch.zeros_like(h_l) if self.predictor_mode == "diagonal": soma = self.Bnuis[l] * h_l else: soma = h_l @ self.Bnuis[l].t() if self.traffic_mode == "soma": traffic = soma else: if context is None: raise ValueError("topdown apical traffic requires a context state") if self.predictor_mode == "diagonal": if context.shape[1] != h_l.shape[1]: raise ValueError( "diagonal topdown traffic requires constant hidden width; " f"got context={context.shape[1]} and layer={h_l.shape[1]}") topdown = self.Bcontext[l] * context else: topdown = context @ self.Bcontext[l].t() traffic = topdown if self.traffic_mode == "topdown" else (soma + topdown) / _SQRT2 return self.nuis_rho * traffic def vectorizer(self, l, c, h_l=None, context=None): """Return the learned instructional component of apical activity. ``soma_gated`` adds ``tanh(h_i) G_i c`` independently in every cell. ``context_gated`` adds a linear readout of ``c outer tanh(context)``. Bounding the state feature prevents a bilinear positive-feedback loop; both terms vanish when c=0, so neutral-period predictor identification remains unchanged. """ value = c @ self.A[l].t() if self.vectorizer_mode == "linear": return value if self.vectorizer_mode == "soma_gated": if h_l is None: raise ValueError("soma-gated vectorizer requires somatic activity") return value + torch.tanh(h_l) * (c @ self.A_gate[l].t()) if context is None: raise ValueError("context-gated vectorizer requires a context state") features = (c.unsqueeze(2) * torch.tanh(context).unsqueeze(1)).flatten(1) return value + features @ self.A_gate[l].t() def apical(self, l, c, h_l=None, context=None): """a_l = A_l c + ordinary non-teaching apical traffic.""" a = self.vectorizer(l, c, h_l, context) if h_l is not None: a = a + self.apical_traffic(l, h_l, context) 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, context=None): """r_l = a_l - ahat_l (or raw a_l if use_residual=False -> ablation).""" a = self.apical(l, c, h_l, context) if use_residual: return a - self.baseline(l, h_l), a return a, a def teaching_signal(net, l, c, h_l, cfg, context=None): """Select the signal that drives a hidden layer's local plasticity. ``match_innovation_norm`` is a deliberately strong raw-apical control: it preserves each sample's raw apical direction exactly while rescaling its Euclidean norm to the norm of the corresponding innovation. Therefore a raw-vs-residual difference under this control cannot be attributed merely to the nuisance making the teaching vector larger. Returns ``(teaching, raw_apical, innovation)`` so training and diagnostic probes use one definition of every signal. """ a = net.apical(l, c, h_l, context) innovation = a - net.baseline(l, h_l) if cfg.use_residual: teaching = innovation elif cfg.raw_scale_control == "match_innovation_norm": raw_norm = a.norm(dim=1, keepdim=True) innovation_norm = innovation.norm(dim=1, keepdim=True) teaching = a * (innovation_norm / raw_norm.clamp_min(1e-12)) else: teaching = a return teaching, a, innovation # -------------------------------------------------------------------------- # 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 simultaneous_node_perturbation_targets(net, x, y, sigma=1e-2, antithetic=True, n_dirs=1): """Estimate every hidden-layer target with shared forward evaluations. Each direction injects an independent Rademacher perturbation at *every* hidden activation during one forward pass. The scalar loss difference is demultiplexed locally as q_l = -dL/dsigma * xi_l. Cross-layer terms add variance but vanish in expectation because the xi_l are independent. This changes causal-calibration cost from layerwise O(n_dirs * depth^2) tail replays to O(n_dirs * depth) full-network work. """ with torch.no_grad(): base = net.forward(x) h = base["h"] batch = x.shape[0] xis = [torch.empty((n_dirs,) + hl.shape, device=hl.device, dtype=hl.dtype) .bernoulli_(0.5).mul_(2).sub_(1) for hl in h[1:-1]] # Fold direction and antithetic sign into the batch dimension. This is # mathematically identical to 2*n_dirs small forwards but turns dozens # of tiny Pascal-GPU launches into one efficient matrix multiplication # per layer. copies = 2 * n_dirs if antithetic else n_dirs x_batch = x.repeat(copies, 1) y_batch = y.repeat(copies) if antithetic: noise_batch = [torch.cat((xi, -xi), dim=0).flatten(0, 1) for xi in xis] else: noise_batch = [xi.flatten(0, 1) for xi in xis] losses = _loss_with_hidden_noise(net, x_batch, y_batch, noise_batch, sigma) losses = losses.view(copies, batch) if antithetic: coeff = (losses[:n_dirs] - losses[n_dirs:]) / (2.0 * sigma) else: coeff = (losses - loss_ce(h[-1], y).unsqueeze(0)) / sigma return [-(coeff.unsqueeze(2) * xi).mean(0) for xi in xis] def _loss_with_hidden_noise(net, x, y, noises, scale): """Forward loss with additive interventions after every hidden layer.""" cur = x for l in range(net.L): pre = cur u = pre @ net.W[l].t() + net.b[l] if l < net.L - 1: act = net.act(u) cur = pre + net.res_alpha * act if net.residual and l >= 1 else act cur = cur + scale * noises[l] else: cur = u return loss_ce(cur, y) 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, pert_mode="layerwise", raw_scale_control="none", direct_node_pert=False, vectorizer_optimizer="sgd", vectorizer_eps=1e-6): self.eta = eta self.eta_A = eta_A self.eta_P = eta_P self.use_residual = use_residual # ablation: residual vs raw apical if raw_scale_control not in ("none", "match_innovation_norm"): raise ValueError(f"unknown raw_scale_control: {raw_scale_control}") self.raw_scale_control = raw_scale_control 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 if pert_mode not in ("layerwise", "simultaneous"): raise ValueError(f"unknown pert_mode: {pert_mode}") self.pert_mode = pert_mode 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 self.direct_node_pert = direct_node_pert if vectorizer_optimizer not in ("sgd", "nlms"): raise ValueError(f"unknown vectorizer optimizer: {vectorizer_optimizer}") self.vectorizer_optimizer = vectorizer_optimizer self.vectorizer_eps = vectorizer_eps def _update_apical_vectorizer(net, h, c, r_list, qs, cfg): """Apply the local causal-regression update to all apical pathways.""" B = c.shape[0] context = h[-2] context_features = None if net.vectorizer_mode == "context_gated": context_features = (c.unsqueeze(2) * torch.tanh(context).unsqueeze(1)).flatten(1) for l in range(net.L - 1): calibration_error = qs[l] - r_list[l] scaled_error = calibration_error if cfg.vectorizer_optimizer == "nlms": if net.vectorizer_mode == "linear": feature_power = c.square().sum(1, keepdim=True) elif net.vectorizer_mode == "soma_gated": # Each cell i sees [c, tanh(h_i)c], so its effective feature # power is ||c||^2 (1+tanh(h_i)^2). feature_power = (c.square().sum(1, keepdim=True) * (1.0 + torch.tanh(h[l + 1]).square())) else: feature_power = (c.square().sum(1, keepdim=True) + context_features.square().sum(1, keepdim=True)) scaled_error = calibration_error / feature_power.clamp_min(cfg.vectorizer_eps) dA = scaled_error.t() @ c / B net.A[l] += cfg.eta_A * dA if net.vectorizer_mode == "soma_gated": dgate = (scaled_error * torch.tanh(h[l + 1])).t() @ c / B net.A_gate[l] += cfg.eta_A * dgate elif net.vectorizer_mode == "context_gated": dgate = scaled_error.t() @ context_features / B net.A_gate[l] += cfg.eta_A * dgate def apical_calibration_step(net, x, y, y_onehot, cfg, pert_mode=None, pert_ndirs=None): """Calibrate the apical vectorizer while keeping all forward weights fixed. This is the feedback-first half of a two-timescale protocol. It uses labels and forward-only causal interventions exactly like an ordinary perturbation event, but performs no hidden or output weight update. The caller controls and accounts for the number of prefix steps and perturbation work. """ if not cfg.learn_A: raise ValueError("apical calibration requires learn_A=True") if cfg.feedback != "error": raise ValueError("feedback-first calibration currently requires error feedback") mode = cfg.pert_mode if pert_mode is None else pert_mode directions = cfg.pert_ndirs if pert_ndirs is None else pert_ndirs if mode not in ("layerwise", "simultaneous"): raise ValueError(f"unknown perturbation mode: {mode}") if directions < 1: raise ValueError("perturbation directions must be positive") with torch.no_grad(): fwd = net.forward(x) h = fwd["h"] logits = h[-1] loss = loss_ce(logits, y).mean().item() c = net.output_error(logits, y_onehot) context = h[-2] r_list = [teaching_signal(net, l, c, h[l + 1], cfg, context)[0] for l in range(net.L - 1)] estimator = (simultaneous_node_perturbation_targets if mode == "simultaneous" else node_perturbation_targets) qs = estimator(net, x, y, sigma=cfg.pert_sigma, n_dirs=directions) _update_apical_vectorizer(net, h, c, r_list, qs, cfg) return loss, {"error": c, "did_pert": True} 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) context = h[-2] # ---- compute innovations for hidden layers ---- r_list, a_list = [], [] for l in range(net.L - 1): r, a, _ = teaching_signal(net, l, c, h[l + 1], cfg, context) 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 or cfg.direct_node_pert) and step % cfg.pert_every == 0) qs = None if did_pert: estimator = (simultaneous_node_perturbation_targets if cfg.pert_mode == "simultaneous" else node_perturbation_targets) qs = estimator(net, x, y, sigma=cfg.pert_sigma, n_dirs=cfg.pert_ndirs) # ================= forward weight updates ================= # hidden layers: three-factor local rule teaching_list = qs if cfg.direct_node_pert else r_list if cfg.direct_node_pert and qs is None: raise RuntimeError("direct node perturbation requires a target on every update step") for l in range(net.L - 1): gain = net.act_prime(u[l]) # phi'(u_l) (B, n_l) delta = teaching_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 and cfg.learn_A: _update_apical_vectorizer(net, h, c, r_list, qs, cfg) # ================= predictor P (neutral) ========================== # KEY identification 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], context) # traffic 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"] context = h[-2] 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, context) # ordinary traffic 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): context = h[-2] for l in range(net.L - 1): r, _, _ = teaching_signal(net, l, c, h[l + 1], cfg, context) 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