diff options
Diffstat (limited to 'experiments')
| -rw-r--r-- | experiments/baseline_smoke.py | 32 |
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") |