chore: capture initial SDIL project state

6 files changed, 1079 insertions, 0 deletions

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)