1 files changed, 131 insertions, 0 deletions
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)