diff options
Diffstat (limited to 'sdil/local_baselines.py')
| -rw-r--r-- | sdil/local_baselines.py | 272 |
1 files changed, 272 insertions, 0 deletions
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 |