diff options
Diffstat (limited to 'sdil/core.py')
| -rw-r--r-- | sdil/core.py | 417 |
1 files changed, 417 insertions, 0 deletions
diff --git a/sdil/core.py b/sdil/core.py new file mode 100644 index 0000000..6137171 --- /dev/null +++ b/sdil/core.py @@ -0,0 +1,417 @@ +""" +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): + # 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)] + # 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)] + + # ---- 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. + 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)] + + # 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: + 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.""" + 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 + q_l = -(dL / sigma^2) * xi (spec convention; scale absorbed by eta_A) + 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 (~ -dL/sigma^2) is not ||grad||, so A learns a mis-scaled + # per-layer step. Normalising each layer's delta to unit RMS decouples + # step size (set by eta) from that miscalibration -- 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) + + # ================= 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) + 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 ======= + did_pert = False + if cfg.learn_A and (step % cfg.pert_every == 0): + did_pert = True + qs = node_perturbation_targets(net, x, y, sigma=cfg.pert_sigma, n_dirs=cfg.pert_ndirs) + 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]) + dP = resid.t() @ h[l + 1] / B # (n_l, n_l) + net.P[l] += cfg.eta_P * dP + + 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) + dP = resid.t() @ hl / B + net.P[l] += eta * dP + + +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 |