feat: add paper-faithful local learning baselines

1 files changed, 107 insertions, 0 deletions

diff --git a/experiments/baseline_smoke.py b/experiments/baseline_smoke.py
new file mode 100644
index 0000000..4bc3d85
--- /dev/null
+++ b/experiments/baseline_smoke.py
@@ -0,0 +1,107 @@
+"""Mechanism checks for the publication baselines (CPU, no dataset needed)."""
+import os
+import sys
+
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from sdil.data import onehot
+from sdil.local_baselines import FANet, PEPITANet, FFNet, EPNet
+
+
+def check_fa_residual_transport():
+    torch.manual_seed(1)
+    x = torch.randn(32, 8)
+    y = torch.randint(0, 3, (32,))
+    yoh = onehot(y, 3)
+    residual = FANet([8, 6, 6, 6, 3], seed=2, residual=True)
+    plain = FANet([8, 6, 6, 6, 3], seed=2, residual=False)
+    for net in (residual, plain):
+        for l in (1, 2):
+            net.B[l].zero_()
+    before_res = [w.clone() for w in residual.W]
+    before_plain = [w.clone() for w in plain.W]
+    residual.fa_step(x, y, yoh, eta=0.01)
+    plain.fa_step(x, y, yoh, eta=0.01)
+    res_changes = [(a - b).norm().item() for a, b in zip(residual.W, before_res)]
+    plain_changes = [(a - b).norm().item() for a, b in zip(plain.W, before_plain)]
+    assert all(v > 0 for v in res_changes)
+    assert plain_changes[0] == 0 and plain_changes[1] == 0
+    assert plain_changes[2] > 0 and plain_changes[3] > 0
+    print("FA residual identity transport:", res_changes)
+
+
+def check_pepita_output_rule():
+    torch.manual_seed(2)
+    net = PEPITANet([5, 4, 3], act="relu", seed=3, keep_prob=1.0)
+    x = torch.rand(16, 5)
+    y = torch.randint(0, 3, (16,))
+    yoh = onehot(y, 3)
+    with torch.no_grad():
+        clean = net._pepita_forward(x, None)
+        error = torch.softmax(clean[-1], 1) - yoh
+        mod = net._pepita_forward(x + error @ net.Fproj.t(), None)
+        mod_error = torch.softmax(mod[-1], 1) - yoh
+        expected = -(mod_error.t() @ mod[-2]) / x.shape[0]
+        old = net.W[-1].clone()
+    net.pepita_step(x, y, yoh, eta=0.1, momentum=0.0)
+    assert torch.allclose(net.W[-1] - old, 0.1 * expected, atol=1e-7, rtol=1e-5)
+    assert net.Fproj.abs().max() <= (6.0 / 5) ** 0.5 * 0.05 + 1e-7
+    print("PEPITA modulated-output delta: exact")
+
+
+def ff_loss(net, layer, x, y, neg):
+    hp = net._inputs_to_layer(layer, net._overlay(x, y))
+    hn = net._inputs_to_layer(layer, net._overlay(x, neg))
+    gp = net._layer_forward(layer, hp).pow(2).mean(1)
+    gn = net._layer_forward(layer, hn).pow(2).mean(1)
+    return (torch.nn.functional.softplus(-gp + net.thr)
+            + torch.nn.functional.softplus(gn - net.thr)).mean().item()
+
+
+def check_ff_local_optimization():
+    torch.manual_seed(3)
+    net = FFNet([20, 32, 32], seed=4)
+    x = torch.randn(128, 20)
+    y = torch.randint(0, 10, (128,))
+    neg = (y + 1) % 10
+    initial = ff_loss(net, 0, x, y, neg)
+    for _ in range(80):
+        net.train_layer(0, x, y, eta=0.03, negative_labels=neg)
+    final = ff_loss(net, 0, x, y, neg)
+    assert final < initial - 0.1
+    assert net._layer_forward(0, x).shape == (128, 32)
+    print(f"FF layer-local loss: {initial:.4f} -> {final:.4f}")
+
+
+def check_ep_energy_dynamics():
+    torch.manual_seed(4)
+    net = EPNet([4, 3, 2], seed=5, dt=0.2, random_beta_sign=False)
+    x = torch.rand(7, 4)
+    y = onehot(torch.randint(0, 2, (7,)), 2)
+    s = [torch.rand(7, 3) * 0.8 + 0.1, torch.rand(7, 2) * 0.8 + 0.1]
+    beta = 0.3
+    manual = []
+    for k in range(net.L):
+        below = x if k == 0 else s[k - 1]
+        drive = -s[k] + below @ net.W[k].t() + net.b[k]
+        if k < net.L - 1:
+            drive = drive + s[k + 1] @ net.W[k + 1]
+        else:
+            drive = drive + 2 * beta * (y - s[k])
+        manual.append((s[k] + net.dt * drive).clamp(0, 1))
+    actual = net._settle(x, y, beta=beta, s=[v.clone() for v in s], T=1)
+    assert all(torch.allclose(a, b, atol=1e-7) for a, b in zip(actual, manual))
+    old = [w.clone() for w in net.W]
+    net.train_step(x, y.argmax(1), y, eta=[0.01, 0.005])
+    assert all(torch.isfinite(w).all() for w in net.W)
+    assert any(not torch.equal(a, b) for a, b in zip(net.W, old))
+    print("EP one-step -d(E+beta*C)/ds dynamics: exact")
+
+
+if __name__ == "__main__":
+    check_fa_residual_transport()
+    check_pepita_output_rule()
+    check_ff_local_optimization()
+    check_ep_energy_dynamics()
+    print("ALL BASELINE SMOKE CHECKS PASSED")