From 6093fe5eb453947154fe39ddb75740f09eb306ef Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Tue, 21 Jul 2026 07:34:09 -0500 Subject: feat: add paper-faithful local learning baselines --- experiments/baseline_run.py | 241 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 241 insertions(+) create mode 100644 experiments/baseline_run.py (limited to 'experiments/baseline_run.py') 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()) -- cgit v1.2.3