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
|
"""Mechanism checks for the publication baselines (CPU, no dataset needed)."""
import os
import sys
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
def check_fa_residual_transport():
torch.manual_seed(1)
x = torch.randn(32, 8)
y = torch.randint(0, 3, (32,))
yoh = onehot(y, 3)
residual = FANet([8, 6, 6, 6, 3], seed=2, residual=True)
plain = FANet([8, 6, 6, 6, 3], seed=2, residual=False)
for net in (residual, plain):
for l in (1, 2):
net.B[l].zero_()
before_res = [w.clone() for w in residual.W]
before_plain = [w.clone() for w in plain.W]
residual.fa_step(x, y, yoh, eta=0.01)
plain.fa_step(x, y, yoh, eta=0.01)
res_changes = [(a - b).norm().item() for a, b in zip(residual.W, before_res)]
plain_changes = [(a - b).norm().item() for a, b in zip(plain.W, before_plain)]
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
print("FA residual identity transport:", res_changes)
def check_pepita_output_rule():
torch.manual_seed(2)
net = PEPITANet([5, 4, 3], act="relu", seed=3, keep_prob=1.0)
x = torch.rand(16, 5)
y = torch.randint(0, 3, (16,))
yoh = onehot(y, 3)
with torch.no_grad():
clean = net._pepita_forward(x, None)
error = torch.softmax(clean[-1], 1) - yoh
mod = net._pepita_forward(x + error @ net.Fproj.t(), None)
mod_error = torch.softmax(mod[-1], 1) - yoh
expected = -(mod_error.t() @ mod[-2]) / x.shape[0]
old = net.W[-1].clone()
net.pepita_step(x, y, yoh, eta=0.1, momentum=0.0)
assert torch.allclose(net.W[-1] - old, 0.1 * expected, atol=1e-7, rtol=1e-5)
assert net.Fproj.abs().max() <= (6.0 / 5) ** 0.5 * 0.05 + 1e-7
print("PEPITA modulated-output delta: exact")
def ff_loss(net, layer, x, y, neg):
hp = net._inputs_to_layer(layer, net._overlay(x, y))
hn = net._inputs_to_layer(layer, net._overlay(x, neg))
gp = net._layer_forward(layer, hp).pow(2).mean(1)
gn = net._layer_forward(layer, hn).pow(2).mean(1)
return (torch.nn.functional.softplus(-gp + net.thr)
+ torch.nn.functional.softplus(gn - net.thr)).mean().item()
def check_ff_local_optimization():
torch.manual_seed(3)
net = FFNet([20, 32, 32], seed=4)
x = torch.randn(128, 20)
y = torch.randint(0, 10, (128,))
neg = (y + 1) % 10
initial = ff_loss(net, 0, x, y, neg)
for _ in range(80):
net.train_layer(0, x, y, eta=0.03, negative_labels=neg)
final = ff_loss(net, 0, x, y, neg)
assert final < initial - 0.1
assert net._layer_forward(0, x).shape == (128, 32)
print(f"FF layer-local loss: {initial:.4f} -> {final:.4f}")
def check_ep_energy_dynamics():
torch.manual_seed(4)
net = EPNet([4, 3, 2], seed=5, dt=0.2, random_beta_sign=False)
x = torch.rand(7, 4)
y = onehot(torch.randint(0, 2, (7,)), 2)
s = [torch.rand(7, 3) * 0.8 + 0.1, torch.rand(7, 2) * 0.8 + 0.1]
beta = 0.3
manual = []
for k in range(net.L):
below = x if k == 0 else s[k - 1]
drive = -s[k] + below @ net.W[k].t() + net.b[k]
if k < net.L - 1:
drive = drive + s[k + 1] @ net.W[k + 1]
else:
drive = drive + 2 * beta * (y - s[k])
manual.append((s[k] + net.dt * drive).clamp(0, 1))
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])
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))
print("EP one-step -d(E+beta*C)/ds dynamics: exact")
if __name__ == "__main__":
check_fa_residual_transport()
check_pepita_output_rule()
check_ff_local_optimization()
check_ep_energy_dynamics()
print("ALL BASELINE SMOKE CHECKS PASSED")
|
