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
132
133
134
135
136
137
138
139
140
141
|
"""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
from sdil import probes
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
alignment = probes.fa_alignment_report(residual, x, y, yoh)["cos_fa_negg"]
assert len(alignment) == 3 and all(torch.isfinite(torch.tensor(alignment)))
print("FA residual identity transport:", res_changes)
print("FA measurable hidden alignment:", alignment)
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))
# Exact author update: contrast weakly-clamped and free correlations,
# divide by beta and minibatch size, then apply the layer-specific rate.
rates = [0.01, 0.005]
old = [w.clone() for w in net.W]
old_biases = [b.clone() for b in net.b]
expected_free = net._settle(x, beta=0.0, T=net.T_free)
expected_nudged = net._settle(
x, y, beta=net.beta, s=[v.clone() for v in expected_free], T=net.T_nudge)
expected_weights, expected_biases = [], []
for k in range(net.L):
below_free = x if k == 0 else net.rho(expected_free[k - 1])
below_nudged = x if k == 0 else net.rho(expected_nudged[k - 1])
correlation_delta = (
net.rho(expected_nudged[k]).t() @ below_nudged
- net.rho(expected_free[k]).t() @ below_free) / (net.beta * x.shape[0])
bias_delta = (net.rho(expected_nudged[k])
- net.rho(expected_free[k])).mean(0) / net.beta
expected_weights.append(old[k] + rates[k] * correlation_delta)
expected_biases.append(old_biases[k] + rates[k] * bias_delta)
_, free_state = net.train_step(x, y.argmax(1), y, eta=rates,
return_free_state=True)
assert all(torch.allclose(actual, expected, atol=1e-7, rtol=1e-5)
for actual, expected in zip(net.W, expected_weights))
assert all(torch.allclose(actual, expected, atol=1e-7, rtol=1e-5)
for actual, expected in zip(net.b, expected_biases))
assert all(torch.equal(actual, expected)
for actual, expected in zip(free_state, expected_free))
continued = net._settle(x, s=free_state, T=1)
restarted = net._settle(x, T=1)
assert any(not torch.equal(a, b) for a, b in zip(continued, restarted))
print("EP one-step -d(E+beta*C)/ds dynamics: exact")
print("EP contrastive parameter update: exact")
print("EP free particles persist across presentations")
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")
|
