1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
|
"""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
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 main():
torch.manual_seed(0)
dev = "cpu"
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"
# ---- 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)
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}"
# ---- CHECK 3: trained-A SDIL beats fixed-random-A (DFA) on alignment ----
print("\nCHECK3 alignment: 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 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)
print(f" residual cos(r,-g) = {sum(al_n['cos_r_negg'])/len(al_n['cos_r_negg']):+.3f}")
print(f" raw apical cos(a,-g) = {sum(al_n['cos_apical_negg'])/len(al_n['cos_apical_negg']):+.3f}")
print(f" pure A c cos(Ac,-g) = {sum(al_n['cos_Ac_negg'])/len(al_n['cos_Ac_negg']):+.3f}")
# ---- 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()
|
