feat: add paper-faithful local learning baselines

3 files changed, 385 insertions, 0 deletions

diff --git a/experiments/baseline_run.py b/experiments/baseline_run.py
new file mode 100644
index 0000000..a96e24d
--- /dev/null
+++ b/experiments/baseline_run.py
@@ -0,0 +1,241 @@
+"""Canonical runners for non-backprop baselines outside the SDIL family.
+
+The methods intentionally do not all share one architecture: Forward-Forward
+has goodness layers instead of a classifier readout, while Equilibrium
+Propagation is an energy-based recurrent network with symmetric connections.
+Forcing either into SDILNet would make the comparison look uniform while
+silently changing the algorithm.  Every JSON result records the protocol.
+"""
+import argparse
+import json
+import os
+import subprocess
+import sys
+import time
+
+import torch
+import torch.nn.functional as F
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from sdil.data import get_dataset, onehot
+from sdil.local_baselines import FANet, PEPITANet, FFNet, EPNet
+
+
+METHOD_SOURCES = {
+    "fa": {
+        "paper": "https://www.nature.com/articles/ncomms13276",
+        "protocol": "fixed random sequential feedback matrices",
+    },
+    "pepita": {
+        "paper": "https://proceedings.mlr.press/v162/dellaferrera22a.html",
+        "code": "https://github.com/GiorgiaD/PEPITA",
+        "protocol": "ERIN two-forward-pass rule; He-uniform; ReLU; F scale 0.05",
+    },
+    "ff": {
+        "paper": "https://www.cs.toronto.edu/~hinton/absps/FFXfinal.pdf",
+        "reference": "https://github.com/mpezeshki/pytorch_forward_forward",
+        "protocol": "greedy layerwise goodness; input length normalization; Adam",
+    },
+    "ep": {
+        "paper": "https://www.frontiersin.org/journals/computational-neuroscience/articles/10.3389/fncom.2017.00024/full",
+        "code": "https://github.com/bscellier/Towards-a-Biologically-Plausible-Backprop",
+        "protocol": "two-phase energy-gradient dynamics; random beta sign",
+    },
+}
+
+_MNIST_STATS = {"mnist": (0.1307, 0.3081), "fmnist": (0.2860, 0.3530)}
+_CIFAR_MEAN = torch.tensor([0.4914, 0.4822, 0.4465])
+_CIFAR_STD = torch.tensor([0.2470, 0.2435, 0.2616])
+
+
+def code_provenance():
+    root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    try:
+        commit = subprocess.run(["git", "rev-parse", "HEAD"], cwd=root, check=True,
+                                capture_output=True, text=True).stdout.strip()
+        dirty = bool(subprocess.run(
+            ["git", "status", "--porcelain", "--untracked-files=no"], cwd=root,
+            check=True, capture_output=True, text=True).stdout.strip())
+        return {"git_commit": commit, "git_dirty": dirty}
+    except (OSError, subprocess.CalledProcessError):
+        return {"git_commit": None, "git_dirty": None}
+
+
+def canonical_input(x, dataset, method, enabled=True):
+    """Restore [0,1] pixels where the authors' protocol requires them."""
+    if not enabled or method not in ("pepita", "ep"):
+        return x
+    if dataset in _MNIST_STATS:
+        mean, std = _MNIST_STATS[dataset]
+        return (x * std + mean).clamp(0, 1)
+    if dataset == "cifar10":
+        mean = _CIFAR_MEAN.to(x.device, x.dtype).view(1, 3, 1)
+        std = _CIFAR_STD.to(x.device, x.dtype).view(1, 3, 1)
+        return (x.view(-1, 3, 1024) * std + mean).clamp(0, 1).view_as(x)
+    return x
+
+
+@torch.no_grad()
+def evaluate_method(net, loader, method, dataset, canonical):
+    correct = total = 0
+    loss_sum = 0.0
+    for x, y in loader:
+        x = canonical_input(x, dataset, method, canonical)
+        if method == "ff":
+            pred = net.predict(x)
+        elif method == "ep":
+            state = net._settle(x, beta=0.0, T=net.T_free)
+            logits = state[-1]
+            pred = logits.argmax(1)
+            loss_sum += F.mse_loss(logits, onehot(y, 10), reduction="sum").item()
+        else:
+            logits = net.logits(x)
+            pred = logits.argmax(1)
+            loss_sum += F.cross_entropy(logits, y, reduction="sum").item()
+        correct += (pred == y).sum().item()
+        total += y.shape[0]
+    return correct / total, loss_sum / total if method != "ff" else float("nan")
+
+
+def ep_learning_rates(depth, base_eta=None):
+    """Layerwise rates from the authors' one/two/three-hidden-layer runs."""
+    canonical = {
+        1: [0.1, 0.05],
+        2: [0.4, 0.1, 0.01],
+        3: [0.128, 0.032, 0.008, 0.002],
+    }
+    if base_eta is not None:
+        return [base_eta / (4 ** l) for l in range(depth + 1)]
+    if depth in canonical:
+        return canonical[depth]
+    # The original implementation already required rapidly shrinking rates as
+    # depth grew. Continue its depth-3 geometric schedule for an explicit,
+    # logged extrapolation rather than pretending there is a canonical one.
+    return [0.128 / (4 ** l) for l in range(depth + 1)]
+
+
+def build(args, n_in, device):
+    sizes = [n_in] + [args.width] * args.depth + [10]
+    if args.method == "fa":
+        return FANet(sizes, act=args.act, device=device, seed=args.seed,
+                     residual=bool(args.residual), b_scale=args.feedback_scale)
+    if args.method == "pepita":
+        return PEPITANet(sizes, act="relu", device=device, seed=args.seed,
+                         f_scale=args.pepita_f_scale, keep_prob=args.keep_prob)
+    if args.method == "ff":
+        # FF has only goodness layers; a 10-unit SDIL-style output is not part
+        # of the original supervised algorithm.
+        return FFNet([n_in] + [args.width] * args.depth, device=device,
+                     seed=args.seed, threshold=args.ff_threshold)
+    if args.method == "ep":
+        return EPNet(sizes, device=device, seed=args.seed, beta=args.ep_beta,
+                     dt=args.ep_dt, T_free=args.ep_free_steps,
+                     T_nudge=args.ep_nudge_steps, random_beta_sign=True)
+    raise ValueError(args.method)
+
+
+def train(args):
+    torch.manual_seed(args.seed)
+    device = args.device
+    train_loader, test_loader, n_in, n_out = get_dataset(
+        args.dataset, args.batch_size, device=device)
+    net = build(args, n_in, device)
+    canonical = bool(args.canonical_preprocess)
+    log = {
+        "args": vars(args),
+        "method_source": METHOD_SOURCES[args.method],
+        "provenance": code_provenance(),
+        "steps": [],
+        "final": {},
+    }
+    t0 = time.time()
+
+    if args.method == "ff":
+        eta = 0.03 if args.eta is None else args.eta
+        for layer in range(args.depth):
+            for epoch in range(args.epochs):
+                batches = 0
+                for x, y in train_loader:
+                    loss = net.train_layer(layer, x, y, eta)
+                    batches += 1
+                    if args.max_batches and batches >= args.max_batches:
+                        break
+                print(f"[{args.tag}] layer {layer} epoch {epoch} local_loss {loss:.4f}",
+                      flush=True)
+            acc, test_loss = evaluate_method(net, test_loader, args.method,
+                                             args.dataset, canonical)
+            log["steps"].append({"layer_end": layer, "test_acc": acc,
+                                 "local_loss": loss})
+            print(f"[{args.tag}] layer {layer} test_acc {acc:.4f}", flush=True)
+    else:
+        eta = args.eta
+        if eta is None:
+            eta = 0.05 if args.method == "fa" else (0.1 if args.method == "pepita" else None)
+        ep_etas = ep_learning_rates(args.depth, eta) if args.method == "ep" else None
+        for epoch in range(args.epochs):
+            if args.method == "pepita" and epoch in (60, 90):
+                eta *= 0.1
+            batches = 0
+            for x, y in train_loader:
+                x = canonical_input(x, args.dataset, args.method, canonical)
+                yoh = onehot(y, n_out, device=device)
+                if args.method == "fa":
+                    loss = net.fa_step(x, y, yoh, eta, args.momentum)
+                elif args.method == "pepita":
+                    loss = net.pepita_step(x, y, yoh, eta, args.momentum)
+                else:
+                    loss = net.train_step(x, y, yoh, ep_etas)
+                batches += 1
+                if args.max_batches and batches >= args.max_batches:
+                    break
+            acc, test_loss = evaluate_method(net, test_loader, args.method,
+                                             args.dataset, canonical)
+            log["steps"].append({"epoch_end": epoch, "test_acc": acc,
+                                 "test_loss": test_loss, "train_loss": loss})
+            print(f"[{args.tag}] epoch {epoch} loss {loss:.4f} test_acc {acc:.4f}",
+                  flush=True)
+
+    acc, test_loss = evaluate_method(net, test_loader, args.method,
+                                     args.dataset, canonical)
+    log["final"] = {"test_acc": acc, "test_loss": test_loss,
+                    "wall_s": time.time() - t0}
+    os.makedirs(args.outdir, exist_ok=True)
+    path = os.path.join(args.outdir, f"{args.tag}.json")
+    with open(path, "w") as f:
+        json.dump(log, f)
+    print(f"[{args.tag}] DONE test_acc={acc:.4f} -> {path}", flush=True)
+    return log
+
+
+def get_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--method", required=True, choices=["fa", "pepita", "ff", "ep"])
+    p.add_argument("--dataset", default="mnist", choices=["mnist", "fmnist", "cifar10"])
+    p.add_argument("--depth", type=int, default=2)
+    p.add_argument("--width", type=int, default=500)
+    p.add_argument("--epochs", type=int, default=10,
+                   help="FF: epochs per greedy layer; other methods: global epochs")
+    p.add_argument("--batch_size", type=int, default=64)
+    p.add_argument("--max_batches", type=int, default=0)
+    p.add_argument("--eta", type=float, default=None)
+    p.add_argument("--momentum", type=float, default=0.9)
+    p.add_argument("--act", default="tanh", choices=["tanh", "gelu", "silu", "relu"])
+    p.add_argument("--residual", type=int, default=0)
+    p.add_argument("--feedback_scale", type=float, default=1.0)
+    p.add_argument("--pepita_f_scale", type=float, default=0.05)
+    p.add_argument("--keep_prob", type=float, default=0.9)
+    p.add_argument("--ff_threshold", type=float, default=2.0)
+    p.add_argument("--ep_beta", type=float, default=0.5)
+    p.add_argument("--ep_dt", type=float, default=0.5)
+    p.add_argument("--ep_free_steps", type=int, default=20)
+    p.add_argument("--ep_nudge_steps", type=int, default=4)
+    p.add_argument("--canonical_preprocess", type=int, default=1)
+    p.add_argument("--seed", type=int, default=0)
+    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    p.add_argument("--outdir", default="results")
+    p.add_argument("--tag", default="baseline")
+    return p.parse_args()
+
+
+if __name__ == "__main__":
+    train(get_args())
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")
diff --git a/experiments/baseline_sweep.sh b/experiments/baseline_sweep.sh
new file mode 100755
index 0000000..bd3ff7b
--- /dev/null
+++ b/experiments/baseline_sweep.sh
@@ -0,0 +1,37 @@
+#!/usr/bin/env bash
+# Run canonical non-backprop baselines. Method-specific knobs can be appended.
+# Usage: baseline_sweep.sh <gpu> "<methods>" "<seeds>" <dataset> <depth> <width> <epochs> <prefix> [extra args...]
+set -eu
+
+cd "$(dirname "$0")/.."
+PY=/home/yurenh2/miniconda3/envs/ep_pascal/bin/python3
+GPU="${1:?GPU index required}"
+METHODS="${2:-fa pepita ff ep}"
+SEEDS="${3:-0}"
+DATASET="${4:-mnist}"
+DEPTH="${5:-2}"
+WIDTH="${6:-500}"
+EPOCHS="${7:-10}"
+PREFIX="${8:-canonical}"
+shift 8 || true
+
+export CUDA_VISIBLE_DEVICES="$GPU"
+export OMP_NUM_THREADS=2
+mkdir -p results logs/baselines
+
+for seed in $SEEDS; do
+  for method in $METHODS; do
+    tag="${PREFIX}_${DATASET}_${method}_w${WIDTH}_d${DEPTH}_s${seed}"
+    result="results/${tag}.json"
+    log="logs/baselines/${tag}.log"
+    if [[ -f "$result" ]]; then
+      echo "skip $tag (result exists)"
+      continue
+    fi
+    echo ">>> $tag $(date --iso-8601=seconds) gpu=$GPU"
+    "$PY" experiments/baseline_run.py --method "$method" --dataset "$DATASET" \
+      --depth "$DEPTH" --width "$WIDTH" --epochs "$EPOCHS" --seed "$seed" \
+      --tag "$tag" --outdir results "$@" > "$log" 2>&1
+    grep -h DONE "$log"
+  done
+done