test: verify exact EP contrastive update

1 files changed, 28 insertions, 4 deletions

diff --git a/experiments/baseline_smoke.py b/experiments/baseline_smoke.py
index 24bb2a4..c656a2f 100644
--- a/experiments/baseline_smoke.py
+++ b/experiments/baseline_smoke.py
@@ -96,16 +96,40 @@ def check_ep_energy_dynamics():
         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]
-    _, free_state = net.train_step(x, y.argmax(1), y, eta=[0.01, 0.005],
+    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.isfinite(w).all() for w in net.W)
-    assert any(not torch.equal(a, b) for a, b in zip(net.W, old))
-    assert len(free_state) == net.L
+    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")