From 215acff518eea2d6d8a5bd7c90c398f49cc3b41b Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Tue, 21 Jul 2026 06:38:55 -0500 Subject: chore: capture initial SDIL project state --- sdil/__init__.py | 1 + sdil/baselines.py | 58 +++++++ sdil/core.py | 417 ++++++++++++++++++++++++++++++++++++++++++++++++ sdil/data.py | 200 +++++++++++++++++++++++ sdil/local_baselines.py | 272 +++++++++++++++++++++++++++++++ sdil/probes.py | 131 +++++++++++++++ 6 files changed, 1079 insertions(+) create mode 100644 sdil/__init__.py create mode 100644 sdil/baselines.py create mode 100644 sdil/core.py create mode 100644 sdil/data.py create mode 100644 sdil/local_baselines.py create mode 100644 sdil/probes.py (limited to 'sdil') diff --git a/sdil/__init__.py b/sdil/__init__.py new file mode 100644 index 0000000..4a7cf44 --- /dev/null +++ b/sdil/__init__.py @@ -0,0 +1 @@ +from . import core, probes, data, baselines # noqa: F401 diff --git a/sdil/baselines.py b/sdil/baselines.py new file mode 100644 index 0000000..f809f83 --- /dev/null +++ b/sdil/baselines.py @@ -0,0 +1,58 @@ +"""Baselines sharing SDILNet's exact architecture and initialisation, so +comparisons are apples-to-apples. + +- BP: exact backprop (autograd) -- the alignment/performance upper bound. +- DFA is NOT a separate class: it is exactly SDIL with a fixed-random apical + vectorizer, no residual, no predictor, no perturbation. See dfa_config(). +""" + +import torch +import torch.nn.functional as F + +from .core import SDILNet, SDILConfig + + +class BPNet(SDILNet): + """Same tanh MLP, trained by exact backprop via autograd.""" + + def bp_step(self, x, y, eta, momentum=0.0): + for W in self.W: + W.requires_grad_(True) + for b in self.b: + b.requires_grad_(True) + logits = self.logits(x) + loss = F.cross_entropy(logits, y) + grads = torch.autograd.grad(loss, self.W + self.b) + nW = len(self.W) + with torch.no_grad(): + for i in range(nW): + if momentum: + self.mW[i].mul_(momentum).add_(grads[i]) + self.mb[i].mul_(momentum).add_(grads[nW + i]) + self.W[i] -= eta * self.mW[i] + self.b[i] -= eta * self.mb[i] + else: + self.W[i] -= eta * grads[i] + self.b[i] -= eta * grads[nW + i] + for W in self.W: + W.requires_grad_(False) + for b in self.b: + b.requires_grad_(False) + return loss.item() + + +def dfa_config(eta=0.05, momentum=0.0): + """SDIL configured to reproduce Direct Feedback Alignment exactly.""" + return SDILConfig(eta=eta, use_residual=False, learn_A=False, learn_P=False, + pert_every=10**9, momentum=momentum) + + +@torch.no_grad() +def evaluate(net, test_loader): + correct, total, tot_loss = 0, 0, 0.0 + for x, y in test_loader: + logits = net.logits(x) + tot_loss += F.cross_entropy(logits, y, reduction="sum").item() + correct += (logits.argmax(1) == y).sum().item() + total += y.shape[0] + return correct / total, tot_loss / total 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 diff --git a/sdil/data.py b/sdil/data.py new file mode 100644 index 0000000..f5b4928 --- /dev/null +++ b/sdil/data.py @@ -0,0 +1,200 @@ +"""Data loaders. Parse the raw IDX files directly (no torchvision dependency, +so it runs on the Pascal env on the 1080 farm). Datasets already live on the +shared NFS home; the box is offline for dataset mirrors.""" + +import pickle +import struct +import numpy as np +import torch + +DATA_DIR = "/home/yurenh2/sdrn/data" + +# (subdir, mean, std) matching the standard torchvision normalisation used by +# the predecessor project, so numbers are comparable. +_STATS = { + "mnist": ("MNIST", 0.1307, 0.3081), + "fmnist": ("FashionMNIST", 0.2860, 0.3530), +} + +_CIFAR_MEAN = np.array([0.4914, 0.4822, 0.4465], dtype=np.float32) +_CIFAR_STD = np.array([0.2470, 0.2435, 0.2616], dtype=np.float32) + + +def _read_images(path): + with open(path, "rb") as f: + magic, n, r, c = struct.unpack(">IIII", f.read(16)) + assert magic == 2051, f"bad image magic {magic}" + buf = f.read(n * r * c) + return np.frombuffer(buf, dtype=np.uint8).reshape(n, r * c).astype(np.float32) + + +def _read_labels(path): + with open(path, "rb") as f: + magic, n = struct.unpack(">II", f.read(8)) + assert magic == 2049, f"bad label magic {magic}" + buf = f.read(n) + return np.frombuffer(buf, dtype=np.uint8).astype(np.int64) + + +def _load_split(subdir, train, mean, std, data_dir): + raw = f"{data_dir}/{subdir}/raw" + pre = "train" if train else "t10k" + x = _read_images(f"{raw}/{pre}-images-idx3-ubyte") + y = _read_labels(f"{raw}/{pre}-labels-idx1-ubyte") + x = (x / 255.0 - mean) / std # normalise + flatten (already 784) + return torch.from_numpy(x), torch.from_numpy(y) + + +class _FastLoader: + """Minimal shuffled minibatch iterator over in-memory (GPU) tensors.""" + def __init__(self, x, y, batch_size, shuffle, seed=0): + self.x, self.y = x, y + self.bs = batch_size + self.shuffle = shuffle + self.n = x.shape[0] + self.g = torch.Generator(device="cpu").manual_seed(seed) + + def __iter__(self): + idx = torch.randperm(self.n, generator=self.g) if self.shuffle else torch.arange(self.n) + for i in range(0, self.n, self.bs): + j = idx[i:i + self.bs] + yield self.x[j], self.y[j] + + def __len__(self): + return (self.n + self.bs - 1) // self.bs + + +def _load_cifar(data_dir, n_classes=10): + root = f"{data_dir}/cifar-10-batches-py" + def unpickle(fn): + with open(fn, "rb") as f: + return pickle.load(f, encoding="bytes") + xs, ys = [], [] + for i in range(1, 6): + d = unpickle(f"{root}/data_batch_{i}") + xs.append(d[b"data"]); ys.append(np.array(d[b"labels"])) + xtr = np.concatenate(xs).astype(np.float32); ytr = np.concatenate(ys).astype(np.int64) + dte = unpickle(f"{root}/test_batch") + xte = np.array(dte[b"data"], dtype=np.float32); yte = np.array(dte[b"labels"], dtype=np.int64) + + def norm(x): # x: (N, 3072) as R(1024)G(1024)B(1024) + x = x.reshape(-1, 3, 1024) / 255.0 + x = (x - _CIFAR_MEAN[None, :, None]) / _CIFAR_STD[None, :, None] + return x.reshape(-1, 3072).astype(np.float32) + return (torch.from_numpy(norm(xtr)), torch.from_numpy(ytr), + torch.from_numpy(norm(xte)), torch.from_numpy(yte)) + + +def get_dataset(name="mnist", batch_size=128, data_dir=DATA_DIR, device="cpu", preload=True): + if name == "cifar10": + xtr, ytr, xte, yte = _load_cifar(data_dir) + n_in = 3072 + else: + subdir, mean, std = _STATS[name] + xtr, ytr = _load_split(subdir, True, mean, std, data_dir) + xte, yte = _load_split(subdir, False, mean, std, data_dir) + n_in = 784 + xtr, ytr = xtr.to(device), ytr.to(device) + xte, yte = xte.to(device), yte.to(device) + return (_FastLoader(xtr, ytr, batch_size, True), + _FastLoader(xte, yte, 1000, False), n_in, 10) + + +def make_teacher_student(n_in=128, n_classes=10, t_depth=8, t_width=64, + n_train=50000, n_test=10000, seed=0, residual=True, + t_scale=1.5, batch_size=128, device="cpu"): + """Deep teacher-student classification. Labels = argmax of a fixed random + deep (residual) teacher net. Students are WIDTH-LIMITED, so depth is the + resource that matters: a shallow student cannot compose enough nonlinearity + to match a deep teacher, a deep one can. This is the standard construction + for a task where accuracy genuinely GROWS with model depth — unlike + flattened-CIFAR MLPs, which are architecture-bottlenecked and depth-flat.""" + from .core import SDILNet + tsizes = [n_in] + [t_width] * t_depth + [n_classes] + teacher = SDILNet(tsizes, act="tanh", device=device, seed=seed + 999, + w_scale=t_scale, residual=residual) + g = torch.Generator(device="cpu").manual_seed(seed) + xtr = torch.randn(n_train, n_in, generator=g).to(device) + xte = torch.randn(n_test, n_in, generator=g).to(device) + with torch.no_grad(): + ytr = teacher.logits(xtr).argmax(1) + yte = teacher.logits(xte).argmax(1) + return (_FastLoader(xtr, ytr, batch_size, True), + _FastLoader(xte, yte, 1000, False), n_in, n_classes) + + +def make_hierarchical(levels=6, n_classes=10, n_train=50000, n_test=10000, + seed=0, batch_size=128, device="cpu"): + """Hierarchical compositional target with PROVABLE depth-dependence. Input + has 2^levels features; C independent binary trees each fold adjacent pairs + with a fixed random nonlinear combiner v' = tanh(a·l + b·r + c·l·r) (the + l·r product is the depth-hard bit). Label = argmax over the C tree roots. + A depth-d net can only realise ~d tree levels, so accuracy GROWS with depth + up to `levels` — the regime where scaling the model actually helps, and where + a rule with poor deep credit assignment (DFA) cannot follow.""" + n_in = 2 ** levels + g = torch.Generator(device="cpu").manual_seed(seed + 31) + # per-class, per-level, per-node combiner params + combiners = [] + width = n_in + for lev in range(levels): + half = width // 2 + combiners.append(torch.randn(n_classes, half, 3, generator=g)) + width = half + + def teacher(x): # x: (N, n_in) + N = x.shape[0] + outs = [] + for cls in range(n_classes): + v = x + w = n_in + for lev in range(levels): + half = w // 2 + l = v[:, 0:2 * half:2] + r = v[:, 1:2 * half:2] + p = combiners[lev][cls].to(x.device) # (half, 3) + v = torch.tanh(l * p[:, 0] + r * p[:, 1] + (l * r) * p[:, 2]) + w = half + outs.append(v[:, 0]) + return torch.stack(outs, 1) # (N, C) + + xtr = torch.randn(n_train, n_in, generator=g).to(device) + xte = torch.randn(n_test, n_in, generator=g).to(device) + with torch.no_grad(): + rtr = teacher(xtr) + rte = teacher(xte) + # standardise per-class roots on train stats -> roughly balanced argmax + mu, sd = rtr.mean(0, keepdim=True), rtr.std(0, keepdim=True) + 1e-6 + ytr = ((rtr - mu) / sd).argmax(1) + yte = ((rte - mu) / sd).argmax(1) + return (_FastLoader(xtr, ytr, batch_size, True), + _FastLoader(xte, yte, 1000, False), n_in, n_classes) + + +def make_tentmap(levels=8, n_in=4, n_train=50000, n_test=10000, seed=0, + batch_size=128, device="cpu"): + """Telgarsky depth-separation task. Label = [tent^levels(x0) > 0.5], where + tent(v)=1-|2v-1| self-composed `levels` times produces 2^levels oscillations. + A depth-d ReLU net can realise ~d tent compositions, resolving only ~2^d of + the 2^levels teeth, so TEST ACCURACY PROVABLY RISES WITH DEPTH up to `levels`. + This is the regime where scaling the model genuinely buys accuracy, and where + a rule with poor deep credit assignment (DFA/FA) cannot build the composition + and stalls. Extra input dims are distractors (must be ignored).""" + g = torch.Generator(device="cpu").manual_seed(seed + 7) + xtr = torch.rand(n_train, n_in, generator=g).to(device) + xte = torch.rand(n_test, n_in, generator=g).to(device) + + def label(x): + v = x[:, 0].clone() + for _ in range(levels): + v = 1.0 - (2.0 * v - 1.0).abs() + return (v > 0.5).long() + + return (_FastLoader(xtr, label(xtr), batch_size, True), + _FastLoader(xte, label(xte), 1000, False), n_in, 2) + + +def onehot(y, n_classes, device=None): + oh = torch.zeros(y.shape[0], n_classes, device=y.device if device is None else device) + oh.scatter_(1, y.view(-1, 1), 1.0) + return oh diff --git a/sdil/local_baselines.py b/sdil/local_baselines.py new file mode 100644 index 0000000..a6a11a3 --- /dev/null +++ b/sdil/local_baselines.py @@ -0,0 +1,272 @@ +""" +Other biologically-motivated / non-backprop local learning baselines, sharing +SDILNet's architecture and initialisation for a fair comparison: + +- FANet : Feedback Alignment (Lillicrap 2016) -- backprop through FIXED random + feedback matrices instead of W^T (sequential, layer-wise). DFA is the + direct variant (already in baselines.dfa_config); FA is the sequential one. +- PEPITANet : PEPITA (Dellaferrera & Kreiman 2022) -- forward-only. A second forward + pass on error-modulated input; weights follow the activation change. +- FFNet : Forward-Forward (Hinton 2022) -- each layer locally maximises a + "goodness" (sum of squares) on positive (real label) data and minimises + it on negative (wrong label) data; inference picks max total goodness. + +All train with no global backward graph. Autograd, where used (FF's per-layer +local loss), never crosses layer boundaries -> the update stays local. +""" +import math +import torch +import torch.nn.functional as F + +from .core import SDILNet, ACTS + + +# -------------------------------------------------------------------------- +# Feedback Alignment (sequential, layer-wise random feedback) +# -------------------------------------------------------------------------- +class FANet(SDILNet): + def __init__(self, *args, b_scale=1.0, **kw): + super().__init__(*args, **kw) + g = torch.Generator(device="cpu").manual_seed(4242) + # fixed random feedback B[l] with the same shape as W[l], for l=1..L-1 + self.B = [None] + for l in range(1, self.L): + shape = self.W[l].shape + self.B.append((torch.randn(*shape, generator=g) * (b_scale / math.sqrt(shape[1])) + ).to(self.W[l].device, self.dtype)) + + def fa_step(self, x, y, yoh, eta, momentum=0.0): + with torch.no_grad(): + fwd = self.forward(x) + h, u = fwd["h"], fwd["u"] + B = x.shape[0] + e = torch.softmax(h[-1], 1) - yoh # grad wrt logits + loss = F.cross_entropy(h[-1], y).item() + # output layer (exact local) + delta = e # grad wrt u[L-1] + self._upd(self.L - 1, delta, h[self.L - 1], eta, momentum) + # hidden layers: propagate with fixed random B + for l in range(self.L - 2, -1, -1): + delta = (delta @ self.B[l + 1]) * self.act_prime(u[l]) # grad wrt u[l] + self._upd(l, delta, h[l], eta, momentum) + return loss + + def _upd(self, l, delta, pre, eta, momentum): + dW = -(delta.t() @ pre) / delta.shape[0] + db = -delta.mean(0) + if momentum: + self.mW[l].mul_(momentum).add_(dW); self.mb[l].mul_(momentum).add_(db) + self.W[l] += eta * self.mW[l]; self.b[l] += eta * self.mb[l] + else: + self.W[l] += eta * dW; self.b[l] += eta * db + + +# -------------------------------------------------------------------------- +# PEPITA (forward-only, error-modulated second pass) +# -------------------------------------------------------------------------- +class PEPITANet(SDILNet): + def __init__(self, *args, f_scale=0.1, **kw): + super().__init__(*args, **kw) + g = torch.Generator(device="cpu").manual_seed(7777) + n_in = self.sizes[0] + # projection of output error back onto the input (fixed random) + self.Fproj = (torch.randn(n_in, self.n_classes, generator=g) + * (f_scale / math.sqrt(self.n_classes))).to(self.W[0].device, self.dtype) + + def pepita_step(self, x, y, yoh, eta, momentum=0.0): + with torch.no_grad(): + std = self.forward(x) + hs = std["h"] + e = torch.softmax(hs[-1], 1) - yoh + loss = F.cross_entropy(hs[-1], y).item() + x_mod = x + (e @ self.Fproj.t()) # modulate the input by the error + mod = self.forward(x_mod) + hm = mod["h"] + B = x.shape[0] + for l in range(self.L): + # change in this layer's output between clean and modulated passes + dpost = hs[l + 1] - hm[l + 1] + # PEPITA: first layer uses the CLEAN input as presynaptic; deeper + # layers use the modulated presynaptic. + pre = x if l == 0 else hm[l] + dW = -(dpost.t() @ pre) / B + db = -dpost.mean(0) + if momentum: + self.mW[l].mul_(momentum).add_(dW); self.mb[l].mul_(momentum).add_(db) + self.W[l] += eta * self.mW[l]; self.b[l] += eta * self.mb[l] + else: + self.W[l] += eta * dW; self.b[l] += eta * db + return loss + + +# -------------------------------------------------------------------------- +# Forward-Forward (Hinton 2022) +# -------------------------------------------------------------------------- +class FFNet: + """Fully-connected Forward-Forward net. Label is overlaid on the input; + each layer is trained by a local goodness objective; classification sums + goodness across layers over the candidate labels.""" + + def __init__(self, sizes, act="tanh", device="cpu", seed=0, threshold=2.0, + n_classes=10, dtype=torch.float32, overlay_val=10.0): + self.sizes = list(sizes) + self.L = len(sizes) - 1 + self.n_classes = n_classes + self.device = device + self.dtype = dtype + self.act, _ = ACTS[act] + self.thr = threshold + self.overlay_val = overlay_val + g = torch.Generator(device="cpu").manual_seed(seed) + self.W, self.b, self.mW, self.mb = [], [], [], [] + for i in range(self.L): + w = (torch.randn(sizes[i + 1], sizes[i], generator=g) / math.sqrt(sizes[i]) + ).to(device, dtype) + w.requires_grad_(True) + self.W.append(w) + bb = torch.zeros(sizes[i + 1], device=device, dtype=dtype, requires_grad=True) + self.b.append(bb) + + def __init_overlay_scale__(self): + pass + + def _overlay(self, x, labels): + """Overlay one-hot label onto the first n_classes input features. + The label value must be strong enough to actually shift the goodness + (10 of 784 pixels is otherwise swamped by the shared image).""" + xo = x.clone() + xo[:, :self.n_classes] = 0.0 + xo[torch.arange(x.shape[0]), labels] = self.overlay_val + return xo + + @staticmethod + def _norm(h): + return h / (h.norm(dim=1, keepdim=True) + 1e-8) + + def _layer_forward(self, l, h_in): + return self.act(h_in @ self.W[l].t() + self.b[l]) + + def train_step(self, x, y, eta): + # positive = correct label; negative = a random wrong label + neg = (y + torch.randint(1, self.n_classes, y.shape, device=y.device)) % self.n_classes + xpos = self._overlay(x, y) + xneg = self._overlay(x, neg) + hpos, hneg = xpos, xneg + total = 0.0 + for l in range(self.L): + hp = self._layer_forward(l, hpos.detach()) + hn = self._layer_forward(l, hneg.detach()) + gp = hp.pow(2).mean(1) # goodness (positive) + gn = hn.pow(2).mean(1) # goodness (negative) + # push positive goodness above threshold, negative below + loss = (F.softplus(-(gp - self.thr)) + F.softplus(gn - self.thr)).mean() + gW, gb = torch.autograd.grad(loss, [self.W[l], self.b[l]]) + with torch.no_grad(): + self.W[l] -= eta * gW + self.b[l] -= eta * gb + total += loss.item() + hpos, hneg = self._norm(hp).detach(), self._norm(hn).detach() + return total / self.L + + @torch.no_grad() + def predict(self, x): + scores = torch.zeros(x.shape[0], self.n_classes, device=x.device) + for c in range(self.n_classes): + h = self._overlay(x, torch.full((x.shape[0],), c, device=x.device, dtype=torch.long)) + good = torch.zeros(x.shape[0], device=x.device) + for l in range(self.L): + h = self._layer_forward(l, h) + if l > 0: # skip first layer for scoring (Hinton) + good = good + h.pow(2).mean(1) + h = self._norm(h) + scores[:, c] = good + return scores.argmax(1) + + @torch.no_grad() + def evaluate(self, test_loader): + correct, total = 0, 0 + for x, y in test_loader: + pred = self.predict(x) + correct += (pred == y).sum().item(); total += y.shape[0] + return correct / total, 0.0 + + +# -------------------------------------------------------------------------- +# Equilibrium Propagation (Scellier & Bengio 2017), real-valued, layered MLP +# -------------------------------------------------------------------------- +class EPNet: + """Prototypical EP: bidirectional layered net, hard-sigmoid rho. Free phase + settles to equilibrium; a weakly-nudged phase perturbs the output toward the + target; the local contrastive update approximates the loss gradient. + Updates are local (products of adjacent-layer rho's) with no backprop.""" + + def __init__(self, sizes, device="cpu", seed=0, beta=0.5, dt=0.5, + T_free=20, T_nudge=8, dtype=torch.float32): + self.sizes = list(sizes) + self.L = len(sizes) - 1 # number of weight layers + self.device = device + self.dtype = dtype + self.beta, self.dt, self.T_free, self.T_nudge = beta, dt, T_free, T_nudge + g = torch.Generator(device="cpu").manual_seed(seed) + self.W = [(torch.randn(sizes[i + 1], sizes[i], generator=g) / math.sqrt(sizes[i]) + ).to(device, dtype) for i in range(self.L)] + self.b = [torch.zeros(sizes[i + 1], device=device, dtype=dtype) for i in range(self.L)] + + @staticmethod + def rho(s): + return s.clamp(0, 1) + + @staticmethod + def rhop(s): + return ((s > 0) & (s < 1)).to(s.dtype) + + def _settle(self, x, y=None, beta=0.0, s=None, T=20): + rx = self.rho(x) + if s is None: + # init in the active region (rhop=0 at the clamp boundaries would + # otherwise freeze the dynamics); a feedforward warm start settles fast. + s = [] + below = rx + for i in range(self.L): + below = (below @ self.W[i].t() + self.b[i]).clamp(0, 1) + s.append(below) + for _ in range(T): + new = [] + for k in range(self.L): + below = rx if k == 0 else self.rho(s[k - 1]) + pre = below @ self.W[k].t() + self.b[k] + if k < self.L - 1: # top-down from layer above + pre = pre + self.rho(s[k + 1]) @ self.W[k + 1] + if k == self.L - 1 and beta: # nudge output toward target + pre = pre + beta * (y - s[k]) + ds = self.rhop(s[k]) * pre - s[k] + new.append((s[k] + self.dt * ds).clamp(0, 1)) + s = new + return s + + def train_step(self, x, y, yoh, eta): + with torch.no_grad(): + s0 = self._settle(x, beta=0.0, T=self.T_free) # free phase + sb = self._settle(x, yoh, beta=self.beta, s=[t.clone() for t in s0], T=self.T_nudge) + B = x.shape[0] + for k in range(self.L): + below0 = self.rho(x) if k == 0 else self.rho(s0[k - 1]) + belowb = self.rho(x) if k == 0 else self.rho(sb[k - 1]) + dW = (self.rho(sb[k]).t() @ belowb - self.rho(s0[k]).t() @ below0) / (self.beta * B) + db = (self.rho(sb[k]) - self.rho(s0[k])).mean(0) / self.beta + self.W[k] += eta * dW + self.b[k] += eta * db + # free-phase output as prediction proxy for loss logging + return F.mse_loss(s0[-1], yoh).item() + + @torch.no_grad() + def predict(self, x): + s = self._settle(x, beta=0.0, T=self.T_free) + return s[-1].argmax(1) + + @torch.no_grad() + def evaluate(self, test_loader): + correct, total = 0, 0 + for x, y in test_loader: + correct += (self.predict(x) == y).sum().item(); total += y.shape[0] + return correct / total, 0.0 diff --git a/sdil/probes.py b/sdil/probes.py new file mode 100644 index 0000000..5489c6e --- /dev/null +++ b/sdil/probes.py @@ -0,0 +1,131 @@ +""" +Measurement probes. These use autograd to compute the TRUE gradient purely for +diagnostics -- the learning rules in core.py never call these. The central +quantity is the alignment cos(r_l, -grad_{h_l}L): SDIL claims the innovation +r_l points along the causal descent direction without weight transport. +""" + +import copy +import torch +import torch.nn.functional as F + +from . import core + + +def true_hidden_grads(net, x, y): + """Return (list of g_l = dL/dh_l over hidden layers, scalar loss). + Uses autograd; temporarily flags W to build the graph, then restores.""" + with torch.enable_grad(): + for W in net.W: + W.requires_grad_(True) + x = x.detach() + h = x + hs = [] + for l in range(net.L): + pre = h @ net.W[l].t() + net.b[l] + if l < net.L - 1: + act = net.act(pre) + h = (h + net.res_alpha * act) if (getattr(net, "residual", False) and l >= 1) else act + hs.append(h) + else: + logits = pre + loss = F.cross_entropy(logits, y) + grads = torch.autograd.grad(loss, hs, retain_graph=False) + for W in net.W: + W.requires_grad_(False) + return [g.detach() for g in grads], loss.item() + + +def _row_cos(u, v, eps=1e-12): + """Mean per-sample cosine between rows of u and v.""" + num = (u * v).sum(1) + den = u.norm(dim=1) * v.norm(dim=1) + eps + return (num / den).mean().item() + + +@torch.no_grad() +def alignment_report(net, x, y, y_onehot, cfg): + """Per-hidden-layer alignments of the teaching signals with -grad. + Returns dict of lists (index = hidden layer).""" + grads, loss = true_hidden_grads(net, x, y) # g_l + fwd = net.forward(x) + h = fwd["h"] + u = fwd["u"] + logits = h[-1] + e = net.output_error(logits, y_onehot) + c = e + + out = {"cos_r_negg": [], "cos_apical_negg": [], "cos_Ac_negg": [], + "r_norm": [], "g_norm": [], "loss": loss} + for l in range(net.L - 1): + g = grads[l] + negg = -g + a = net.apical(l, c, h[l + 1]) # raw apical (with nuisance) + Ac = c @ net.A[l].t() # pure vectorizer output + r = a - net.baseline(l, h[l + 1]) # innovation + # include the phi' gain that the plasticity rule actually applies, + # so the reported alignment is what the weight update "sees" + gain = net.act_prime(u[l]) + out["cos_r_negg"].append(_row_cos(r * gain, negg * gain)) + out["cos_apical_negg"].append(_row_cos(a * gain, negg * gain)) + out["cos_Ac_negg"].append(_row_cos(Ac * gain, negg * gain)) + out["r_norm"].append(r.norm(dim=1).mean().item()) + out["g_norm"].append(g.norm(dim=1).mean().item()) + return out + + +@torch.no_grad() +def loss_decrease_ratio(net, x, y, y_onehot, cfg, step): + """Single-step descent quality: apply one SDIL update to a scratch copy and + measure the actual loss drop on the same batch, compared to one plain-SGD + (true-gradient) step at the same learning rate. Ratio in (0,1] means SDIL + descends but less steeply than exact GD; <=0 means it went uphill.""" + y_ = y + L0 = core.loss_ce(net.logits(x), y_).mean().item() + + # SDIL step on a copy + net_sdil = _clone(net) + core.sdil_step(net_sdil, x, y_, y_onehot, cfg, step) + L_sdil = core.loss_ce(net_sdil.logits(x), y_).mean().item() + + # exact-gradient SGD step on a copy at the same eta + net_bp = _clone(net) + _sgd_step(net_bp, x, y_, cfg.eta) + L_bp = core.loss_ce(net_bp.logits(x), y_).mean().item() + + d_sdil = L_sdil - L0 + d_bp = L_bp - L0 + ratio = d_sdil / d_bp if abs(d_bp) > 1e-9 else float("nan") + return {"dL_sdil": d_sdil, "dL_bp": d_bp, "ratio": ratio} + + +def _clone(net): + n = copy.copy(net) + n.W = [w.clone() for w in net.W] + 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.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] + return n + + +def _sgd_step(net, x, y, eta): + with torch.enable_grad(): + for W in net.W: + W.requires_grad_(True) + for b in net.b: + b.requires_grad_(True) + logits = net.logits(x) + loss = F.cross_entropy(logits, y) + grads = torch.autograd.grad(loss, net.W + net.b) + nW = len(net.W) + with torch.no_grad(): + for i in range(nW): + net.W[i] -= eta * grads[i] + net.b[i] -= eta * grads[nW + i] + for W in net.W: + W.requires_grad_(False) + for b in net.b: + b.requires_grad_(False) -- cgit v1.2.3