"""Fast CPU correctness checks for SDIL mechanics and signs. Run: python experiments/smoke.py """ import os import sys import torch sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from sdil.core import (SDILNet, SDILConfig, apical_calibration_step, sdil_step, node_perturbation_targets, simultaneous_node_perturbation_targets, neutral_p_update, teaching_signal, _update_apical_vectorizer) from sdil.baselines import dfa_config from sdil import probes from sdil.data import get_dataset, onehot def row_cos(u, v, eps=1e-12): return ((u * v).sum(1) / (u.norm(dim=1) * v.norm(dim=1) + eps)).mean().item() def flat_cos(u, v, eps=1e-12): return (u.flatten() @ v.flatten() / (u.norm() * v.norm() + eps)).item() def check_residual_local_jacobian(): """The three-factor rule must include the residual branch's alpha.""" torch.manual_seed(7) batch = 32 net = SDILNet([20, 16, 16, 16, 5], device="cpu", seed=4, residual=True) x = torch.randn(batch, 20) y = torch.randint(0, 5, (batch,)) for weight in net.W: weight.requires_grad_(True) fwd = net.forward(x) loss = torch.nn.functional.cross_entropy(fwd["h"][-1], y) hidden_grads = torch.autograd.grad(loss, fwd["h"][1:-1], retain_graph=True) weight_grads = torch.autograd.grad(loss, net.W) for weight in net.W: weight.requires_grad_(False) print("CHECK0 residual local-rule Jacobian:") for l in range(net.L - 1): # autograd's hidden gradient includes the mean over the batch; recover # per-example errors before applying the same minibatch mean as SDIL. error = -hidden_grads[l] * batch delta = error * net.act_prime(fwd["u"][l]) if l >= 1: delta = net.res_alpha * delta local_update = delta.t() @ fwd["h"][l] / batch exact_update = -weight_grads[l] cosine = flat_cos(local_update, exact_update) norm_ratio = (local_update.norm() / exact_update.norm()).item() print(f" layer {l}: cos={cosine:.6f} norm_ratio={norm_ratio:.6f}") assert cosine > 0.99999 assert abs(norm_ratio - 1.0) < 1e-5 def check_neutral_predictor(): """Per-neuron neutral regression should remove normal dendritic coupling.""" torch.manual_seed(11) net = SDILNet([20, 32, 32, 5], device="cpu", seed=5, nuis_rho=2.0, predictor_mode="diagonal") x = torch.randn(128, 20) initial = [] final = [] with torch.no_grad(): h = net.forward(x)["h"] zero_c = torch.zeros(x.shape[0], net.n_classes) for l in range(net.L - 1): target = net.apical(l, zero_c, h[l + 1]) initial.append((target - net.baseline(l, h[l + 1])).norm() / target.norm()) for _ in range(300): neutral_p_update(net, x, 0.05) with torch.no_grad(): h = net.forward(x)["h"] for l in range(net.L - 1): target = net.apical(l, zero_c, h[l + 1]) final.append((target - net.baseline(l, h[l + 1])).norm() / target.norm()) print("CHECK0b neutral predictor residual/target:") for l, (before, after) in enumerate(zip(initial, final)): print(f" layer {l}: {before.item():.4f} -> {after.item():.4f}") assert after < 0.03 def check_topdown_predictor(): """A local soma predictor should remove the predictable part of feedback generated by the network's own high-level contextual state.""" torch.manual_seed(13) net = SDILNet([20, 32, 32, 32, 5], device="cpu", seed=6, nuis_rho=1.0, predictor_mode="diagonal", residual=True, traffic_mode="topdown") x = torch.randn(256, 20) def unexplained_fraction(): with torch.no_grad(): h = net.forward(x)["h"] context = h[-2] fractions = [] for l in range(net.L - 1): traffic = net.apical_traffic(l, h[l + 1], context) residual = traffic - net.baseline(l, h[l + 1]) fractions.append((residual.pow(2).mean() / traffic.pow(2).mean()).item()) return fractions initial = unexplained_fraction() for _ in range(500): neutral_p_update(net, x, 0.02) final = unexplained_fraction() print("CHECK0c top-down traffic unexplained power:") for l, (before, after) in enumerate(zip(initial, final)): print(f" layer {l}: {before:.4f} -> {after:.4f}") assert after < before * 0.8 assert final[-1] < 0.03 def check_traffic_seed_isolation(): """Traffic-family seeds must not change feedforward initialization/data.""" net_a = SDILNet([20, 32, 32, 32, 5], device="cpu", seed=6, nuis_rho=1.0, nuis_seed=41, residual=True, traffic_mode="topdown") net_b = SDILNet([20, 32, 32, 32, 5], device="cpu", seed=6, nuis_rho=1.0, nuis_seed=42, residual=True, traffic_mode="topdown") x = torch.randn(64, 20) for wa, wb in zip(net_a.W, net_b.W): assert torch.equal(wa, wb) ha = net_a.forward(x)["h"] hb = net_b.forward(x)["h"] for xa, xb in zip(ha, hb): assert torch.equal(xa, xb) traffic_a = net_a.apical_traffic(0, ha[1], ha[-2]) traffic_b = net_b.apical_traffic(0, hb[1], hb[-2]) assert not torch.equal(traffic_a, traffic_b) print("CHECK0d traffic seed changes feedback only: passed") def check_state_conditioned_vectorizers(): """Gated apical features must start as linear feedback and learn locally.""" torch.manual_seed(17) x = torch.randn(48, 6) y = torch.randint(0, 2, (48,)) yoh = onehot(y, 2) for mode in ("soma_gated", "context_gated"): net = SDILNet([6, 8, 8, 2], act="relu", device="cpu", seed=9, residual=True, vectorizer_mode=mode) fwd = net.forward(x) c = net.output_error(fwd["h"][-1], yoh) context = fwd["h"][-2] before = net.vectorizer(0, c, fwd["h"][1], context) assert torch.allclose(before, c @ net.A[0].t()) gates_before = [gate.clone() for gate in net.A_gate] cfg = SDILConfig(eta=0.01, eta_A=0.02, pert_every=1, pert_ndirs=2, pert_mode="simultaneous") sdil_step(net, x, y, yoh, cfg, step=0) assert any(not torch.equal(old, new) for old, new in zip(gates_before, net.A_gate)) print("CHECK0e state-conditioned vectorizers: zero-init and local calibration passed") def check_direct_node_perturbation(): """Unamortized q must update hidden weights without using A.""" torch.manual_seed(19) x = torch.randn(32, 5) y = torch.randint(0, 3, (32,)) yoh = onehot(y, 3) net = SDILNet([5, 7, 7, 3], act="tanh", device="cpu", seed=10, residual=True) for weight in net.A: weight.zero_() weights_before = [weight.clone() for weight in net.W] apical_before = [weight.clone() for weight in net.A] rng_state = torch.get_rng_state() fwd = net.forward(x) targets = simultaneous_node_perturbation_targets( net, x, y, sigma=0.01, n_dirs=2) expected = [] for layer, target in enumerate(targets): delta = target * net.act_prime(fwd["u"][layer]) if layer >= 1: delta = net.res_alpha * delta expected.append(delta.t() @ fwd["h"][layer] / x.shape[0]) torch.set_rng_state(rng_state) cfg = SDILConfig(eta=0.01, learn_A=False, learn_P=False, pert_every=1, pert_ndirs=2, pert_mode="simultaneous", direct_node_pert=True) sdil_step(net, x, y, yoh, cfg, step=0) for layer in range(net.L - 1): observed = net.W[layer] - weights_before[layer] assert torch.allclose(observed, cfg.eta * expected[layer], atol=1e-6, rtol=1e-5) assert all(torch.equal(before, after) for before, after in zip(apical_before, net.A)) print("CHECK0f direct node perturbation: exact q update; A untouched") def check_feedback_first_calibration(): """A-only calibration must leave the entire forward network unchanged.""" torch.manual_seed(23) x = torch.randn(32, 5) y = torch.randint(0, 3, (32,)) yoh = onehot(y, 3) net = SDILNet([5, 7, 7, 3], act="relu", device="cpu", seed=11, residual=True, vectorizer_mode="context_gated") weights_before = [weight.clone() for weight in net.W] biases_before = [bias.clone() for bias in net.b] apical_before = [weight.clone() for weight in net.A] gates_before = [weight.clone() for weight in net.A_gate] cfg = SDILConfig(eta_A=0.02, learn_A=True, learn_P=True, pert_ndirs=1, pert_mode="simultaneous") apical_calibration_step( net, x, y, yoh, cfg, pert_mode="simultaneous", pert_ndirs=2) assert all(torch.equal(before, after) for before, after in zip(weights_before, net.W)) assert all(torch.equal(before, after) for before, after in zip(biases_before, net.b)) assert any(not torch.equal(before, after) for before, after in zip(apical_before, net.A)) assert any(not torch.equal(before, after) for before, after in zip(gates_before, net.A_gate)) print("CHECK0g feedback-first calibration: A changed; W/readout frozen") def check_vectorizer_nlms(): """NLMS must divide the joint linear/context update by feature power.""" torch.manual_seed(29) batch = 16 net = SDILNet([5, 7, 7, 3], act="relu", device="cpu", seed=12, residual=True, vectorizer_mode="context_gated") for weight in net.A: weight.zero_() for weight in net.A_gate: weight.zero_() h = net.forward(torch.randn(batch, 5))["h"] c = torch.randn(batch, 3) qs = [torch.randn_like(hidden) for hidden in h[1:-1]] residuals = [torch.zeros_like(target) for target in qs] cfg = SDILConfig(eta_A=0.03, vectorizer_optimizer="nlms", vectorizer_eps=1e-6) features = (c.unsqueeze(2) * torch.tanh(h[-2]).unsqueeze(1)).flatten(1) denominator = (c.square().sum(1, keepdim=True) + features.square().sum(1, keepdim=True)).clamp_min(cfg.vectorizer_eps) expected_a0 = cfg.eta_A * ((qs[0] / denominator).t() @ c / batch) expected_gate0 = cfg.eta_A * ((qs[0] / denominator).t() @ features / batch) _update_apical_vectorizer(net, h, c, residuals, qs, cfg) assert torch.allclose(net.A[0], expected_a0, atol=1e-7, rtol=1e-6) assert torch.allclose(net.A_gate[0], expected_gate0, atol=1e-7, rtol=1e-6) assert torch.isfinite(net.A[0]).all() and torch.isfinite(net.A_gate[0]).all() print("CHECK0h vectorizer NLMS: exact joint-feature normalization") def main(): torch.manual_seed(0) dev = "cpu" check_residual_local_jacobian() check_neutral_predictor() check_topdown_predictor() check_traffic_seed_isolation() check_state_conditioned_vectorizers() check_direct_node_perturbation() check_feedback_first_calibration() check_vectorizer_nlms() print("loading MNIST subset...") tr, te, n_in, n_out = get_dataset("mnist", batch_size=128, device=dev) xb, yb = next(iter(tr)) x = xb[:256] y = yb[:256] yoh = onehot(y, 10) sizes = [784, 64, 64, 64, 10] # ---- CHECK 1: node-perturbation q estimates the descent direction -g ---- net = SDILNet(sizes, device=dev, seed=1) grads, _ = probes.true_hidden_grads(net, x, y) for ndirs in (1, 8, 32): qs = node_perturbation_targets(net, x, y, sigma=1e-2, n_dirs=ndirs) cs = [row_cos(qs[l], -grads[l]) for l in range(len(qs))] print(f"CHECK1 node-pert n_dirs={ndirs:2d} cos(q,-g) per layer = " + " ".join(f"{c:+.3f}" for c in cs)) assert all(c > 0 for c in cs), "node perturbation q should align with -grad" qs = simultaneous_node_perturbation_targets(net, x, y, sigma=1e-2, n_dirs=32) cs = [row_cos(qs[l], -grads[l]) for l in range(len(qs))] print("CHECK1 simultaneous n_dirs=32 cos(q,-g) per layer = " + " ".join(f"{c:+.3f}" for c in cs)) assert all(c > 0.2 for c in cs), "simultaneous perturbation should align in expectation" # ---- CHECK 2: SDIL overfits a fixed batch and alignment climbs ---- net = SDILNet(sizes, device=dev, seed=2) cfg = SDILConfig(eta=0.1, eta_A=0.05, eta_P=0.005, pert_every=2, pert_ndirs=8, momentum=0.9) initial_alignment = probes.alignment_report(net, x, y, yoh, cfg) initial_mean_cos = sum(initial_alignment["cos_r_negg"]) / len(initial_alignment["cos_r_negg"]) print("\nCHECK2 SDIL overfit fixed batch:") for step in range(401): loss, _ = sdil_step(net, x, y, yoh, cfg, step) if step % 50 == 0: al = probes.alignment_report(net, x, y, yoh, cfg) mc = sum(al["cos_r_negg"]) / len(al["cos_r_negg"]) mca = sum(al["cos_apical_negg"]) / len(al["cos_apical_negg"]) print(f" step {step:3d} loss {loss:.4f} mean_cos(r,-g) {mc:+.3f} " f"mean_cos(apical,-g) {mca:+.3f} per-layer_r {['%.2f'%v for v in al['cos_r_negg']]}") assert loss < 1.5, f"SDIL should reduce loss on fixed batch, got {loss}" final_mean_cos = sum(al["cos_r_negg"]) / len(al["cos_r_negg"]) assert final_mean_cos > initial_mean_cos + 0.1, ( f"trained A should improve alignment: {initial_mean_cos:.3f} -> {final_mean_cos:.3f}") # ---- CHECK 3: report DFA as a sanity comparator. On a single repeatedly # trained batch, ordinary feedback alignment can exceed learned A, so no # universal ordering is asserted here; CHECK2 tests that A actually learns. print("\nCHECK3 alignment sanity comparison: SDIL(trained A) vs DFA(fixed A):") net_dfa = SDILNet(sizes, device=dev, seed=3) cfg_dfa = dfa_config(eta=0.1, momentum=0.9) for step in range(401): sdil_step(net_dfa, x, y, yoh, cfg_dfa, step) al_dfa = probes.alignment_report(net_dfa, x, y, yoh, cfg_dfa) al_sdil = probes.alignment_report(net, x, y, yoh, cfg) print(f" DFA mean_cos(r,-g) = {sum(al_dfa['cos_r_negg'])/len(al_dfa['cos_r_negg']):+.3f}") print(f" SDIL mean_cos(r,-g) = {sum(al_sdil['cos_r_negg'])/len(al_sdil['cos_r_negg']):+.3f}") # ---- CHECK 4: nuisance -> residual preserves alignment, raw apical degrades ---- print("\nCHECK4 residualization under soma-predictable nuisance (rho=2.0):") net_n = SDILNet(sizes, device=dev, seed=4, nuis_rho=2.0) cfg_n = SDILConfig(eta=0.1, eta_A=0.05, eta_P=0.02, pert_every=2, pert_ndirs=8, momentum=0.9) for _ in range(200): neutral_p_update(net_n, x, 0.05) for step in range(401): sdil_step(net_n, x, y, yoh, cfg_n, step) al_n = probes.alignment_report(net_n, x, y, yoh, cfg_n) residual_cos = sum(al_n['cos_r_negg']) / len(al_n['cos_r_negg']) apical_cos = sum(al_n['cos_apical_negg']) / len(al_n['cos_apical_negg']) print(f" residual cos(r,-g) = {residual_cos:+.3f}") print(f" raw apical cos(a,-g) = {apical_cos:+.3f}") print(f" pure A c cos(Ac,-g) = {sum(al_n['cos_Ac_negg'])/len(al_n['cos_Ac_negg']):+.3f}") assert residual_cos > apical_cos + 0.1 cfg_raw = SDILConfig(use_residual=False, learn_A=False, learn_P=True) al_raw = probes.alignment_report(net_n, x, y, yoh, cfg_raw) assert al_raw["cos_r_negg"] == al_raw["cos_apical_negg"], ( "reported teaching alignment must honor use_residual=False") # A magnitude-matched raw signal is still pointed along raw apical activity, # but has exactly the innovation's norm for every sample. This isolates the # directional effect of subtracting the soma-predictable component. cfg_matched = SDILConfig(use_residual=False, learn_A=False, learn_P=True, raw_scale_control="match_innovation_norm") with torch.no_grad(): fwd_n = net_n.forward(x) error_n = net_n.output_error(fwd_n["h"][-1], yoh) for l in range(net_n.L - 1): matched, raw, innovation = teaching_signal( net_n, l, error_n, fwd_n["h"][l + 1], cfg_matched) assert torch.allclose(matched.norm(dim=1), innovation.norm(dim=1), atol=1e-6, rtol=1e-5) assert row_cos(matched, raw) > 0.99999 al_matched = probes.alignment_report(net_n, x, y, yoh, cfg_matched) for matched_cos, raw_cos in zip(al_matched["cos_r_negg"], al_matched["cos_apical_negg"]): assert abs(matched_cos - raw_cos) < 1e-6 print(" matched raw: direction preserved; per-sample innovation norm matched") # ---- CHECK 5: single-step loss-decrease ratio vs exact GD ---- print("\nCHECK5 single-step loss-decrease ratio vs exact GD:") ldr = probes.loss_decrease_ratio(net, x, y, yoh, cfg, step=100) print(f" dL_sdil={ldr['dL_sdil']:+.4f} dL_bp={ldr['dL_bp']:+.4f} ratio={ldr['ratio']:+.3f}") print("\nALL SMOKE CHECKS PASSED") if __name__ == "__main__": main()