diff options
| -rw-r--r-- | experiments/baseline_run.py | 136 | ||||
| -rw-r--r-- | experiments/baseline_smoke.py | 12 | ||||
| -rwxr-xr-x | experiments/ep_original_sweep.sh | 32 | ||||
| -rw-r--r-- | sdil/data.py | 9 | ||||
| -rw-r--r-- | sdil/local_baselines.py | 46 | ||||
| -rw-r--r-- | sdil/probes.py | 18 |
6 files changed, 215 insertions, 38 deletions
diff --git a/experiments/baseline_run.py b/experiments/baseline_run.py index a96e24d..24ad6ae 100644 --- a/experiments/baseline_run.py +++ b/experiments/baseline_run.py @@ -19,6 +19,7 @@ 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 +from sdil import probes METHOD_SOURCES = { @@ -39,7 +40,7 @@ METHOD_SOURCES = { "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", + "protocol": "two-phase energy-gradient dynamics; persistent free particles; random beta sign", }, } @@ -108,10 +109,54 @@ def ep_learning_rates(depth, base_eta=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)] + raise ValueError("non-canonical EP depth requires an explicit --eta") + + +def ep_dynamics(depth, beta=None, free_steps=None, nudge_steps=None): + """Original-model dynamics for the published 1/2/3-hidden-layer runs.""" + canonical = { + 1: {"beta": 0.5, "free_steps": 20, "nudge_steps": 4}, + 2: {"beta": 1.0, "free_steps": 150, "nudge_steps": 6}, + 3: {"beta": 1.0, "free_steps": 500, "nudge_steps": 8}, + } + if depth not in canonical and (beta is None or free_steps is None or nudge_steps is None): + raise ValueError( + "non-canonical EP depth requires --ep_beta, --ep_free_steps, and --ep_nudge_steps") + defaults = canonical.get(depth, {}) + return { + "beta": defaults["beta"] if beta is None else beta, + "free_steps": defaults["free_steps"] if free_steps is None else free_steps, + "nudge_steps": defaults["nudge_steps"] if nudge_steps is None else nudge_steps, + "canonical_depth": depth in canonical, + } + + +def pepita_protocol(args): + """Resolve only settings that are actually supported by a cited protocol. + + The original PEPITA paper specifies one hidden layer. Deep PEPITA is very + sensitive to architecture-specific learning rates, so silently reusing the + shallow 0.1 default is worse than requiring an explicit value. + """ + eta = args.eta + if eta is None: + if args.depth != 1: + raise ValueError("deep PEPITA requires an explicit --eta") + if args.dataset == "mnist": + eta = 0.1 + elif args.dataset == "cifar10": + eta = 0.01 + else: + raise ValueError("PEPITA on this dataset requires an explicit --eta") + if args.pepita_decay_epochs is not None: + decay_epochs = [int(v) for v in args.pepita_decay_epochs.split(",") if v.strip()] + elif args.depth == 1 and args.dataset == "mnist": + decay_epochs = [60] + elif args.depth == 1 and args.dataset == "cifar10": + decay_epochs = [60, 90] + else: + decay_epochs = [] + return {"eta": eta, "decay_epochs": decay_epochs} def build(args, n_in, device): @@ -128,30 +173,53 @@ def build(args, n_in, device): 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) + dyn = ep_dynamics(args.depth, args.ep_beta, args.ep_free_steps, + args.ep_nudge_steps) + return EPNet(sizes, device=device, seed=args.seed, beta=dyn["beta"], + dt=args.ep_dt, T_free=dyn["free_steps"], + T_nudge=dyn["nudge_steps"], random_beta_sign=True) raise ValueError(args.method) def train(args): torch.manual_seed(args.seed) device = args.device + pepita_cfg = pepita_protocol(args) if args.method == "pepita" else None + ep_dyn = (ep_dynamics(args.depth, args.ep_beta, args.ep_free_steps, + args.ep_nudge_steps) if args.method == "ep" else None) + use_persistent_ep = args.method == "ep" and bool(args.ep_persistent) train_loader, test_loader, n_in, n_out = get_dataset( - args.dataset, args.batch_size, device=device) + args.dataset, args.batch_size, device=device, + shuffle_train=not use_persistent_ep, + train_limit=args.train_examples or None) net = build(args, n_in, device) canonical = bool(args.canonical_preprocess) + resolved_protocol = {} + if pepita_cfg is not None: + resolved_protocol.update(pepita_cfg) + if ep_dyn is not None: + resolved_protocol.update(ep_dyn) + resolved_protocol["learning_rates"] = ep_learning_rates(args.depth, args.eta) + resolved_protocol["persistent_particles"] = use_persistent_ep + resolved_protocol["train_examples"] = train_loader.n log = { "args": vars(args), + "resolved_protocol": resolved_protocol, "method_source": METHOD_SOURCES[args.method], "provenance": code_provenance(), "steps": [], "final": {}, } + if args.method == "fa": + px, py = next(iter(test_loader)) + px, py = px[:args.probe_bs], py[:args.probe_bs] + pyoh = onehot(py, n_out, device=device) t0 = time.time() if args.method == "ff": eta = 0.03 if args.eta is None else args.eta + resolved_protocol["eta"] = eta + resolved_protocol["epochs_per_layer"] = args.epochs for layer in range(args.depth): for epoch in range(args.epochs): batches = 0 @@ -168,15 +236,21 @@ def train(args): "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) + if args.method == "fa": + eta = 0.05 if args.eta is None else args.eta + resolved_protocol["eta"] = eta + resolved_protocol["feedback_scale"] = args.feedback_scale + elif args.method == "pepita": + eta = pepita_cfg["eta"] + else: + eta = args.eta ep_etas = ep_learning_rates(args.depth, eta) if args.method == "ep" else None + ep_states = [None] * len(train_loader) if use_persistent_ep else None for epoch in range(args.epochs): - if args.method == "pepita" and epoch in (60, 90): + if args.method == "pepita" and epoch in pepita_cfg["decay_epochs"]: eta *= 0.1 batches = 0 - for x, y in train_loader: + for batch_index, (x, y) in enumerate(train_loader): x = canonical_input(x, args.dataset, args.method, canonical) yoh = onehot(y, n_out, device=device) if args.method == "fa": @@ -184,21 +258,33 @@ def train(args): elif args.method == "pepita": loss = net.pepita_step(x, y, yoh, eta, args.momentum) else: - loss = net.train_step(x, y, yoh, ep_etas) + if use_persistent_ep: + loss, ep_states[batch_index] = net.train_step( + x, y, yoh, ep_etas, free_state=ep_states[batch_index], + return_free_state=True) + 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) + rec = {"epoch_end": epoch, "test_acc": acc, + "test_loss": test_loss, "train_loss": loss} + msg = f"[{args.tag}] epoch {epoch} loss {loss:.4f} test_acc {acc:.4f}" + if args.method == "fa": + al = probes.fa_alignment_report(net, px, py, pyoh) + rec["cos_fa_negg"] = al["cos_fa_negg"] + msg += f" mean_cos(fa,-g) {sum(al['cos_fa_negg']) / len(al['cos_fa_negg']):+.3f}" + log["steps"].append(rec) + print(msg, 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} + if args.method == "fa": + log["final"].update(probes.fa_alignment_report(net, px, py, pyoh)) os.makedirs(args.outdir, exist_ok=True) path = os.path.join(args.outdir, f"{args.tag}.json") with open(path, "w") as f: @@ -216,6 +302,8 @@ def get_args(): 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("--train_examples", type=int, default=0, + help="0 uses the full training split") 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) @@ -223,13 +311,17 @@ def get_args(): 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("--pepita_decay_epochs", default=None, + help="comma-separated; deep PEPITA defaults to no decay") 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_beta", type=float, default=None) 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("--ep_free_steps", type=int, default=None) + p.add_argument("--ep_nudge_steps", type=int, default=None) + p.add_argument("--ep_persistent", type=int, default=1) p.add_argument("--canonical_preprocess", type=int, default=1) + p.add_argument("--probe_bs", type=int, default=512) 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") diff --git a/experiments/baseline_smoke.py b/experiments/baseline_smoke.py index 4bc3d85..24bb2a4 100644 --- a/experiments/baseline_smoke.py +++ b/experiments/baseline_smoke.py @@ -7,6 +7,7 @@ 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 +from sdil import probes def check_fa_residual_transport(): @@ -28,7 +29,10 @@ def check_fa_residual_transport(): 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 + alignment = probes.fa_alignment_report(residual, x, y, yoh)["cos_fa_negg"] + assert len(alignment) == 3 and all(torch.isfinite(torch.tensor(alignment))) print("FA residual identity transport:", res_changes) + print("FA measurable hidden alignment:", alignment) def check_pepita_output_rule(): @@ -93,10 +97,16 @@ def check_ep_energy_dynamics(): 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]) + _, free_state = net.train_step(x, y.argmax(1), y, eta=[0.01, 0.005], + return_free_state=True) 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)) + assert len(free_state) == net.L + continued = net._settle(x, s=free_state, T=1) + restarted = net._settle(x, T=1) + assert any(not torch.equal(a, b) for a, b in zip(continued, restarted)) print("EP one-step -d(E+beta*C)/ds dynamics: exact") + print("EP free particles persist across presentations") if __name__ == "__main__": diff --git a/experiments/ep_original_sweep.sh b/experiments/ep_original_sweep.sh new file mode 100755 index 0000000..9bb5947 --- /dev/null +++ b/experiments/ep_original_sweep.sh @@ -0,0 +1,32 @@ +#!/usr/bin/env bash +# Scellier & Bengio (2017) one-hidden-layer MNIST protocol, adapted only to +# evaluate on the standard 10k test split. Persistent particles use the fixed +# first 50k training examples, matching the author's training-set size. +# Usage: ep_original_sweep.sh <gpu> "<seeds>" [prefix] +set -eu + +cd "$(dirname "$0")/.." +PY=/home/yurenh2/miniconda3/envs/ep_pascal/bin/python3 +GPU="${1:?GPU index required}" +SEEDS="${2:-0}" +PREFIX="${3:-ep_original_v2}" + +export CUDA_VISIBLE_DEVICES="$GPU" +export OMP_NUM_THREADS=2 +mkdir -p results logs/baselines + +for seed in $SEEDS; do + tag="${PREFIX}_mnist_ep_w500_d1_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 ep --dataset mnist --depth 1 --width 500 --epochs 25 \ + --batch_size 20 --train_examples 50000 --ep_persistent 1 \ + --seed "$seed" --tag "$tag" --outdir results > "$log" 2>&1 + grep -h DONE "$log" +done diff --git a/sdil/data.py b/sdil/data.py index f5b4928..8253d84 100644 --- a/sdil/data.py +++ b/sdil/data.py @@ -85,7 +85,8 @@ def _load_cifar(data_dir, n_classes=10): torch.from_numpy(norm(xte)), torch.from_numpy(yte)) -def get_dataset(name="mnist", batch_size=128, data_dir=DATA_DIR, device="cpu", preload=True): +def get_dataset(name="mnist", batch_size=128, data_dir=DATA_DIR, device="cpu", preload=True, + shuffle_train=True, train_limit=None): if name == "cifar10": xtr, ytr, xte, yte = _load_cifar(data_dir) n_in = 3072 @@ -94,9 +95,13 @@ def get_dataset(name="mnist", batch_size=128, data_dir=DATA_DIR, device="cpu", p xtr, ytr = _load_split(subdir, True, mean, std, data_dir) xte, yte = _load_split(subdir, False, mean, std, data_dir) n_in = 784 + if train_limit is not None: + if train_limit <= 0 or train_limit > xtr.shape[0]: + raise ValueError(f"train_limit must be in [1, {xtr.shape[0]}], got {train_limit}") + xtr, ytr = xtr[:train_limit], ytr[:train_limit] xtr, ytr = xtr.to(device), ytr.to(device) xte, yte = xte.to(device), yte.to(device) - return (_FastLoader(xtr, ytr, batch_size, True), + return (_FastLoader(xtr, ytr, batch_size, shuffle_train), _FastLoader(xte, yte, 1000, False), n_in, 10) diff --git a/sdil/local_baselines.py b/sdil/local_baselines.py index 23b1b5c..dc33fea 100644 --- a/sdil/local_baselines.py +++ b/sdil/local_baselines.py @@ -41,26 +41,41 @@ class FANet(SDILNet): self.B.append((torch.randn(*shape, generator=g) * (b_scale / math.sqrt(shape[1])) ).to(self.W[l].device, self.dtype)) + def fa_hidden_gradients(self, fwd, yoh): + """Return FA estimates of dL/dh for each hidden block output. + + The returned tensors include the residual-branch scale used by the + corresponding local weight update. Identity skips transport ``dh`` + exactly; only the nonlinear branch is transported through fixed B. + Keeping this calculation separate makes the signal directly + measurable without using autograd for learning. + """ + h, u = fwd["h"], fwd["u"] + e = torch.softmax(h[-1], 1) - yoh + dh = e @ self.B[self.L - 1] + estimates = [None] * (self.L - 1) + for l in range(self.L - 2, -1, -1): + branch_scale = self.res_alpha if self.residual and l >= 1 else 1.0 + estimates[l] = branch_scale * dh + du = estimates[l] * self.act_prime(u[l]) + if l > 0: + feedback = du @ self.B[l] + dh = dh + feedback if self.residual and l >= 1 else feedback + return estimates + def fa_step(self, x, y, yoh, eta, momentum=0.0): with torch.no_grad(): fwd = self.forward(x) h, u = fwd["h"], fwd["u"] e = torch.softmax(h[-1], 1) - yoh # grad wrt logits loss = F.cross_entropy(h[-1], y).item() + hidden_grad = self.fa_hidden_gradients(fwd, yoh) # output layer (exact local) self._upd(self.L - 1, e, h[self.L - 1], eta, momentum) - # dh is the error wrt the hidden *state*. On a residual block, - # h'=h+alpha*phi(Wh), the identity branch transports dh exactly and - # only the nonlinear branch uses a fixed random feedback matrix. - dh = e @ self.B[self.L - 1] - for l in range(self.L - 2, -1, -1): - branch_scale = self.res_alpha if self.residual and l >= 1 else 1.0 - du = branch_scale * dh * self.act_prime(u[l]) + for l in range(self.L - 1): + du = hidden_grad[l] * self.act_prime(u[l]) self._upd(l, du, h[l], eta, momentum) - if l > 0: - feedback = du @ self.B[l] - dh = dh + feedback if self.residual and l >= 1 else feedback return loss def _upd(self, l, delta, pre, eta, momentum): @@ -321,9 +336,13 @@ class EPNet: s = new return s - def train_step(self, x, y, yoh, eta): + def train_step(self, x, y, yoh, eta, free_state=None, return_free_state=False): with torch.no_grad(): - s0 = self._settle(x, beta=0.0, T=self.T_free) # free phase + s0 = self._settle(x, beta=0.0, s=free_state, T=self.T_free) # free phase + # The authors keep one free-phase particle per training example + # and reuse it on the next epoch. Save the free state, not the + # subsequently nudged state. + persistent_state = [t.detach().clone() for t in s0] beta = self.beta if self.random_beta_sign and torch.randint(0, 2, ()).item() == 0: beta = -beta @@ -338,7 +357,8 @@ class EPNet: self.W[k] += layer_eta * dW self.b[k] += layer_eta * db # free-phase output as prediction proxy for loss logging - return F.mse_loss(s0[-1], yoh).item() + loss = F.mse_loss(s0[-1], yoh).item() + return (loss, persistent_state) if return_free_state else loss @torch.no_grad() def predict(self, x): diff --git a/sdil/probes.py b/sdil/probes.py index fad82ca..2d10aad 100644 --- a/sdil/probes.py +++ b/sdil/probes.py @@ -85,6 +85,24 @@ def alignment_report(net, x, y, y_onehot, cfg): @torch.no_grad() +def fa_alignment_report(net, x, y, y_onehot): + """Alignment of sequential-FA local update signals with exact descent. + + ``FANet.fa_hidden_gradients`` performs the same fixed-feedback transport + used during learning. Autograd appears only here to obtain the reference + hidden-state gradients for measurement. + """ + grads, loss = true_hidden_grads(net, x, y) + fwd = net.forward(x) + estimated = net.fa_hidden_gradients(fwd, y_onehot) + cosines = [] + for l, (estimate, true_grad) in enumerate(zip(estimated, grads)): + gain = net.act_prime(fwd["u"][l]) + cosines.append(_row_cos(-estimate * gain, -true_grad * gain)) + return {"cos_fa_negg": cosines, "loss": loss} + + +@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 |