path: root/sdil/core.py

"""
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,
                 predictor_mode="diagonal"):
        # 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)]
        if predictor_mode not in ("diagonal", "full"):
            raise ValueError(f"unknown predictor_mode: {predictor_mode}")
        self.predictor_mode = predictor_mode

        # ---- soma-predictable nuisance on the apical compartment ----
        # Harnett et al. define each cell's SD residual using a least-squares line
        # between that same cell's somatic and dendritic event magnitudes.  The
        # faithful default is therefore a per-neuron affine predictor, not a dense
        # population map.  "full" preserves the original experimental ablation.
        self.nuis_rho = nuis_rho
        gn = torch.Generator(device="cpu").manual_seed(nuis_seed)
        if predictor_mode == "diagonal":
            # Positive log-normal slopes model ordinary soma-dendrite coupling.
            self.Bnuis = [torch.exp(0.25 * torch.randn(sizes[i + 1], generator=gn)).to(
                device=device, dtype=dtype) for i in range(self.L - 1)]
            self.P = [torch.zeros(sizes[i + 1], device=device, dtype=dtype)
                      for i in range(self.L - 1)]
            self.P_bias = [torch.zeros(sizes[i + 1], device=device, dtype=dtype)
                           for i in range(self.L - 1)]
        else:
            self.Bnuis = [he((sizes[i + 1], sizes[i + 1]), 1.0, gen=gn)
                          for i in range(self.L - 1)]
            self.P = [torch.zeros((sizes[i + 1], sizes[i + 1]), device=device, dtype=dtype)
                      for i in range(self.L - 1)]
            self.P_bias = None

        # velocity buffers for momentum on forward weights (optional)
        self.mW = [torch.zeros_like(w) for w in self.W]
        self.mb = [torch.zeros_like(bb) for bb in self.b]

    # ------- forward -------
    def forward(self, x, settle_steps=0, kappa=0.0, c_override=None):
        """Return dict of activations.

        h[0]=x; for l in 1..L-1 hidden tanh; logits = W_L h_{L-1} + b_L (linear).
        If settle_steps>0, run online apical control: after the feedforward pass,
        iterate h_l <- h_l + (1/tau)(-h_l + phi(u_l) + kappa r_l) using the current
        innovation r_l (desired-velocity interpretation). Requires labels via
        c_override or is applied post-hoc by the caller.
        """
        h = [x]
        u = []
        for l in range(self.L):
            pre = h[-1]
            ul = pre @ self.W[l].t() + self.b[l]
            u.append(ul)
            if l < self.L - 1:
                act = self.act(ul)
                # skip on interior blocks only (l>=1 => width->width, dims match)
                if self.residual and l >= 1:
                    h.append(pre + self.res_alpha * act)
                else:
                    h.append(act)
            else:
                h.append(ul)                 # logits (linear output)
        return {"h": h, "u": u}

    def logits(self, x):
        return self.forward(x)["h"][-1]

    # ------- feedback signal -------
    def output_error(self, logits, y_onehot):
        # softmax-CE gradient wrt logits
        p = torch.softmax(logits, dim=1)
        return p - y_onehot

    def apical(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:
            if self.predictor_mode == "diagonal":
                a = a + self.nuis_rho * self.Bnuis[l] * h_l
            else:
                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."""
        if self.predictor_mode == "diagonal":
            return self.P[l] * h_l + self.P_bias[l]
        return h_l @ self.P[l].t()

    def innovation(self, l, c, h_l, use_residual=True):
        """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 the antithetic estimate
        q_l = -((L(+sigma*xi) - L(-sigma*xi)) / (2*sigma)) * xi
    Antithetic pairs (+xi, -xi) cancel the O(sigma^2) even term. Averaging over
    n_dirs independent directions cuts the estimator variance ~1/n_dirs, which
    the predecessor project's failures flagged as the main risk of learning a
    feedback pathway from causal probes.
    """
    with torch.no_grad():
        base = net.forward(x)
        h = base["h"]
        L0 = loss_ce(h[-1], y)                     # (B,)
        qs = []
        for l in range(net.L - 1):                 # hidden layers
            hl = h[l + 1]
            acc = torch.zeros_like(hl)
            for _ in range(n_dirs):
                xi = torch.randn_like(hl)
                dL = _relayer_loss(net, h, l + 1, hl + sigma * xi, y) - L0
                # Correctly-scaled zeroth-order estimator of grad_{h_l}L: with
                # perturbation delta=sigma*xi, ghat = delta*dL/sigma^2 = xi*dL/sigma,
                # so E[ghat]=grad. Dividing by sigma (NOT sigma^2) makes ||q||~||grad||
                # -> the three-factor update is genuine SGD at learning rate eta.
                if antithetic:
                    dL2 = _relayer_loss(net, h, l + 1, hl - sigma * xi, y) - L0
                    coeff = (dL - dL2) / (2.0 * sigma)
                else:
                    coeff = dL / sigma
                acc += -(coeff.unsqueeze(1)) * xi
            qs.append(acc / n_dirs)                 # (B, n_l)
        return qs


def simultaneous_node_perturbation_targets(net, x, y, sigma=1e-2,
                                           antithetic=True, n_dirs=1):
    """Estimate every hidden-layer target with shared forward evaluations.

    Each direction injects an independent Rademacher perturbation at *every*
    hidden activation during one forward pass.  The scalar loss difference is
    demultiplexed locally as q_l = -dL/dsigma * xi_l.  Cross-layer terms add
    variance but vanish in expectation because the xi_l are independent.  This
    changes causal-calibration cost from layerwise O(n_dirs * depth^2) tail
    replays to O(n_dirs * depth) full-network work.
    """
    with torch.no_grad():
        base = net.forward(x)
        h = base["h"]
        batch = x.shape[0]
        xis = [torch.empty((n_dirs,) + hl.shape, device=hl.device, dtype=hl.dtype)
                .bernoulli_(0.5).mul_(2).sub_(1) for hl in h[1:-1]]

        # Fold direction and antithetic sign into the batch dimension.  This is
        # mathematically identical to 2*n_dirs small forwards but turns dozens
        # of tiny Pascal-GPU launches into one efficient matrix multiplication
        # per layer.
        copies = 2 * n_dirs if antithetic else n_dirs
        x_batch = x.repeat(copies, 1)
        y_batch = y.repeat(copies)
        if antithetic:
            noise_batch = [torch.cat((xi, -xi), dim=0).flatten(0, 1) for xi in xis]
        else:
            noise_batch = [xi.flatten(0, 1) for xi in xis]
        losses = _loss_with_hidden_noise(net, x_batch, y_batch, noise_batch, sigma)
        losses = losses.view(copies, batch)
        if antithetic:
            coeff = (losses[:n_dirs] - losses[n_dirs:]) / (2.0 * sigma)
        else:
            coeff = (losses - loss_ce(h[-1], y).unsqueeze(0)) / sigma
        return [-(coeff.unsqueeze(2) * xi).mean(0) for xi in xis]


def _loss_with_hidden_noise(net, x, y, noises, scale):
    """Forward loss with additive interventions after every hidden layer."""
    cur = x
    for l in range(net.L):
        pre = cur
        u = pre @ net.W[l].t() + net.b[l]
        if l < net.L - 1:
            act = net.act(u)
            cur = pre + net.res_alpha * act if net.residual and l >= 1 else act
            cur = cur + scale * noises[l]
        else:
            cur = u
    return loss_ce(cur, y)


def _relayer_loss(net, h, start_idx, h_start_new, y):
    """Recompute loss when hidden activation at layer index `start_idx`
    (1-based into h) is replaced by h_start_new, propagating forward.
    start_idx>=1 so every layer in the tail is an interior/readout layer."""
    cur = h_start_new
    for l in range(start_idx, net.L):
        pre = cur
        ul = pre @ net.W[l].t() + net.b[l]
        if l < net.L - 1:
            act = net.act(ul)
            cur = pre + net.res_alpha * act if net.residual else act    # interior blocks skip (l>=1)
        else:
            cur = ul
    return loss_ce(cur, y)


# --------------------------------------------------------------------------
# the SDIL update
# --------------------------------------------------------------------------
class SDILConfig:
    def __init__(self, eta=0.05, eta_A=0.02, eta_P=0.002,
                 use_residual=True, learn_A=True, learn_P=True,
                 pert_sigma=1e-2, pert_every=5, pert_ndirs=1, momentum=0.0, wd=0.0,
                 settle_steps=0, kappa=0.0, feedback="error", p_update_on_neutral=True,
                 normalize_delta=False, pert_mode="layerwise"):
        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
        if pert_mode not in ("layerwise", "simultaneous"):
            raise ValueError(f"unknown pert_mode: {pert_mode}")
        self.pert_mode = pert_mode
        self.momentum = momentum
        self.wd = wd
        self.settle_steps = settle_steps      # online apical control iterations
        self.kappa = kappa
        self.feedback = feedback              # "error" | "error_deriv"
        self.p_update_on_neutral = p_update_on_neutral
        # node perturbation estimates the descent DIRECTION well but its
        # magnitude is noisy, so A can learn a mis-scaled per-layer step.
        # Normalising each layer's delta to unit RMS decouples step size (set by
        # eta) from that noise -- like normalized SGD.
        self.normalize_delta = normalize_delta


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)

        # Calibrate A against the same network state that produced r and c.
        # Measuring q after mutating W would pair a post-update causal target
        # with a pre-update prediction, introducing an avoidable stale-target
        # error (especially at large learning rates).
        did_pert = cfg.learn_A and (step % cfg.pert_every == 0)
        qs = None
        if did_pert:
            estimator = (simultaneous_node_perturbation_targets
                         if cfg.pert_mode == "simultaneous"
                         else node_perturbation_targets)
            qs = estimator(net, x, y, sigma=cfg.pert_sigma, n_dirs=cfg.pert_ndirs)

        # ================= forward weight updates =================
        # hidden layers: three-factor local rule
        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)
            # For an interior residual block h'=h+alpha*phi(Wh), the local
            # Jacobian dh'/dW contains alpha.  Omitting it preserves direction
            # but makes the effective step grow as 1/alpha=sqrt(depth), which
            # confounds depth-scaling comparisons.  The input projection (l=0)
            # is plain rather than residual and therefore has scale 1.
            if net.residual and l >= 1:
                delta = net.res_alpha * delta
            dW = delta.t() @ h[l] / B                  # (n_l, n_{l-1})
            db = delta.mean(0)
            _apply(net, l, dW, db, cfg)
        # output layer: exact local delta rule (output error is locally available)
        gL = e                                         # grad wrt logits
        dWL = -(gL.t() @ h[net.L - 1]) / B             # descent
        dbL = -gL.mean(0)
        _apply(net, net.L - 1, dWL, dbL, cfg)

        # ================= apical vectorizer A via node perturbation =======
        if did_pert:
            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])
                _update_predictor(net, l, resid, h[l + 1], cfg.eta_P)

    return loss, {"error": e, "did_pert": did_pert}


def neutral_p_update(net, x, eta):
    """Fit the predictor P to the neutral-period (c=0) apical drive, i.e. the
    soma-predictable nuisance. Used both as a pre-task warmup and (optionally)
    ongoing. Because the nuisance is linear in h, a converged P cancels it for
    any h, so warming P up before task plasticity prevents the nuisance from
    corrupting the forward weights early. No autograd."""
    with torch.no_grad():
        fwd = net.forward(x)
        h = fwd["h"]
        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)
            _update_predictor(net, l, resid, hl, eta)


def _update_predictor(net, l, resid, h_l, eta):
    """Local least-squares update for the selected soma-dendrite predictor."""
    if net.predictor_mode == "diagonal":
        # Normalized LMS makes the per-cell slope learning rate independent of
        # that cell's activity variance. Centering separates slope and intercept
        # updates, matching an online least-squares line fit.
        h_centered = h_l - h_l.mean(0)
        resid_centered = resid - resid.mean(0)
        variance = h_centered.pow(2).mean(0)
        net.P[l] += eta * ((resid_centered * h_centered).mean(0) / (variance + 1e-6))
        net.P_bias[l] += eta * resid.mean(0)
    else:
        net.P[l] += eta * (resid.t() @ h_l / h_l.shape[0])


def _settle(net, h, u, c, cfg):
    """Online apical control: tau dh = -h + phi(u) + kappa r, a few Euler steps.
    Updates hidden states in place-ish (returns new h list)."""
    h = list(h)
    dt_over_tau = 1.0 / max(1, cfg.settle_steps)
    for _ in range(cfg.settle_steps):
        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