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
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
|
"""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
from experiments.import_native_baseline import (summarize_burst_history,
summarize_dualprop_hist)
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")
def check_native_result_selection():
burst_rows = [
{"epoch": 0, "epoch/top1_error/test": 30.0,
"epoch/top1_error/val": 28.0, "_runtime": 10.0},
{"epoch": 1, "epoch/top1_error/test": 20.0,
"epoch/top1_error/val": 22.0, "_runtime": 20.0},
{"epoch": 2, "epoch/top1_error/test": 21.0,
"epoch/top1_error/val": 18.0, "_runtime": 30.0},
]
burst = summarize_burst_history(burst_rows, 3)
assert burst["primary_final_test_accuracy_percent"] == 79.0
assert burst["validation_selected_epoch"] == 2
assert burst["validation_selected_test_accuracy_percent"] == 79.0
assert burst["test_selected_epoch"] == 1
assert burst["test_selected_test_accuracy_percent"] == 80.0
dual = summarize_dualprop_hist({
"train_loss": [2.0, 1.0], "train_accuracy": [30.0, 70.0],
"train_time": [4.0, 5.0], "val_loss": [1.5, 1.2],
"val_accuracy": [60.0, 65.0], "val_top5accuracy": [90.0, 92.0],
"val_time": [1.0, 1.5], "test_loss": 1.1,
"test_accuracy": 64.0, "test_top5accuracy": 93.0, "test_time": 2.0,
}, 2)
assert dual["validation_selected_epoch"] == 2
assert dual["primary_test_accuracy_percent"] == 64.0
assert dual["test_evaluations"] == 1
assert dual["wall_s"] == 13.5
print("native final/validation-selected/test-selected semantics: exact")
if __name__ == "__main__":
check_fa_residual_transport()
check_pepita_output_rule()
check_ff_local_optimization()
check_ep_energy_dynamics()
check_native_result_selection()
print("ALL BASELINE SMOKE CHECKS PASSED")
|
