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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
|
"""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 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 main():
torch.manual_seed(0)
dev = "cpu"
check_residual_local_jacobian()
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)
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 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()
|
