path: root/experiments/smoke.py

"""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, sdil_step,
                       node_perturbation_targets, simultaneous_node_perturbation_targets,
                       neutral_p_update, teaching_signal)
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 main():
    torch.manual_seed(0)
    dev = "cpu"
    check_residual_local_jacobian()
    check_neutral_predictor()
    check_topdown_predictor()
    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()