sdil: normalize causal vectorizer regression

1 files changed, 29 insertions, 1 deletions

diff --git a/experiments/smoke.py b/experiments/smoke.py
index 9971800..37c38e3 100644
--- a/experiments/smoke.py
+++ b/experiments/smoke.py
@@ -8,7 +8,7 @@ import torch
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 from sdil.core import (SDILNet, SDILConfig, apical_calibration_step, sdil_step,
                        node_perturbation_targets, simultaneous_node_perturbation_targets,
-                       neutral_p_update, teaching_signal)
+                       neutral_p_update, teaching_signal, _update_apical_vectorizer)
 from sdil.baselines import dfa_config
 from sdil import probes
 from sdil.data import get_dataset, onehot
@@ -215,6 +215,33 @@ def check_feedback_first_calibration():
     print("CHECK0g feedback-first calibration: A changed; W/readout frozen")
 
 
+def check_vectorizer_nlms():
+    """NLMS must divide the joint linear/context update by feature power."""
+    torch.manual_seed(29)
+    batch = 16
+    net = SDILNet([5, 7, 7, 3], act="relu", device="cpu", seed=12,
+                  residual=True, vectorizer_mode="context_gated")
+    for weight in net.A:
+        weight.zero_()
+    for weight in net.A_gate:
+        weight.zero_()
+    h = net.forward(torch.randn(batch, 5))["h"]
+    c = torch.randn(batch, 3)
+    qs = [torch.randn_like(hidden) for hidden in h[1:-1]]
+    residuals = [torch.zeros_like(target) for target in qs]
+    cfg = SDILConfig(eta_A=0.03, vectorizer_optimizer="nlms", vectorizer_eps=1e-6)
+    features = (c.unsqueeze(2) * torch.tanh(h[-2]).unsqueeze(1)).flatten(1)
+    denominator = (c.square().sum(1, keepdim=True)
+                   + features.square().sum(1, keepdim=True)).clamp_min(cfg.vectorizer_eps)
+    expected_a0 = cfg.eta_A * ((qs[0] / denominator).t() @ c / batch)
+    expected_gate0 = cfg.eta_A * ((qs[0] / denominator).t() @ features / batch)
+    _update_apical_vectorizer(net, h, c, residuals, qs, cfg)
+    assert torch.allclose(net.A[0], expected_a0, atol=1e-7, rtol=1e-6)
+    assert torch.allclose(net.A_gate[0], expected_gate0, atol=1e-7, rtol=1e-6)
+    assert torch.isfinite(net.A[0]).all() and torch.isfinite(net.A_gate[0]).all()
+    print("CHECK0h vectorizer NLMS: exact joint-feature normalization")
+
+
 def main():
     torch.manual_seed(0)
     dev = "cpu"
@@ -225,6 +252,7 @@ def main():
     check_state_conditioned_vectorizers()
     check_direct_node_perturbation()
     check_feedback_first_calibration()
+    check_vectorizer_nlms()
     print("loading MNIST subset...")
     tr, te, n_in, n_out = get_dataset("mnist", batch_size=128, device=dev)
     xb, yb = next(iter(tr))