fix: scale local updates by residual Jacobian

2 files changed, 59 insertions, 7 deletions

diff --git a/experiments/smoke.py b/experiments/smoke.py
index f2efa87..092e9f5 100644
--- a/experiments/smoke.py
+++ b/experiments/smoke.py
@@ -16,9 +16,47 @@ def row_cos(u, v, eps=1e-12):
     return ((u * v).sum(1) / (u.norm(dim=1) * v.norm(dim=1) + eps)).mean().item()
 
 
+def flat_cos(u, v, eps=1e-12):
+    return (u.flatten() @ v.flatten() / (u.norm() * v.norm() + eps)).item()
+
+
+def check_residual_local_jacobian():
+    """The three-factor rule must include the residual branch's alpha."""
+    torch.manual_seed(7)
+    batch = 32
+    net = SDILNet([20, 16, 16, 16, 5], device="cpu", seed=4, residual=True)
+    x = torch.randn(batch, 20)
+    y = torch.randint(0, 5, (batch,))
+    for weight in net.W:
+        weight.requires_grad_(True)
+    fwd = net.forward(x)
+    loss = torch.nn.functional.cross_entropy(fwd["h"][-1], y)
+    hidden_grads = torch.autograd.grad(loss, fwd["h"][1:-1], retain_graph=True)
+    weight_grads = torch.autograd.grad(loss, net.W)
+    for weight in net.W:
+        weight.requires_grad_(False)
+
+    print("CHECK0 residual local-rule Jacobian:")
+    for l in range(net.L - 1):
+        # autograd's hidden gradient includes the mean over the batch; recover
+        # per-example errors before applying the same minibatch mean as SDIL.
+        error = -hidden_grads[l] * batch
+        delta = error * net.act_prime(fwd["u"][l])
+        if l >= 1:
+            delta = net.res_alpha * delta
+        local_update = delta.t() @ fwd["h"][l] / batch
+        exact_update = -weight_grads[l]
+        cosine = flat_cos(local_update, exact_update)
+        norm_ratio = (local_update.norm() / exact_update.norm()).item()
+        print(f"  layer {l}: cos={cosine:.6f} norm_ratio={norm_ratio:.6f}")
+        assert cosine > 0.99999
+        assert abs(norm_ratio - 1.0) < 1e-5
+
+
 def main():
     torch.manual_seed(0)
     dev = "cpu"
+    check_residual_local_jacobian()
     print("loading MNIST subset...")
     tr, te, n_in, n_out = get_dataset("mnist", batch_size=128, device=dev)
     xb, yb = next(iter(tr))
@@ -42,6 +80,8 @@ def main():
     net = SDILNet(sizes, device=dev, seed=2)
     cfg = SDILConfig(eta=0.1, eta_A=0.05, eta_P=0.005, pert_every=2, pert_ndirs=8,
                      momentum=0.9)
+    initial_alignment = probes.alignment_report(net, x, y, yoh, cfg)
+    initial_mean_cos = sum(initial_alignment["cos_r_negg"]) / len(initial_alignment["cos_r_negg"])
     print("\nCHECK2 SDIL overfit fixed batch:")
     for step in range(401):
         loss, _ = sdil_step(net, x, y, yoh, cfg, step)
@@ -52,9 +92,14 @@ def main():
             print(f"  step {step:3d} loss {loss:.4f} mean_cos(r,-g) {mc:+.3f} "
                   f"mean_cos(apical,-g) {mca:+.3f} per-layer_r {['%.2f'%v for v in al['cos_r_negg']]}")
     assert loss < 1.5, f"SDIL should reduce loss on fixed batch, got {loss}"
+    final_mean_cos = sum(al["cos_r_negg"]) / len(al["cos_r_negg"])
+    assert final_mean_cos > initial_mean_cos + 0.1, (
+        f"trained A should improve alignment: {initial_mean_cos:.3f} -> {final_mean_cos:.3f}")
 
-    # ---- CHECK 3: trained-A SDIL beats fixed-random-A (DFA) on alignment ----
-    print("\nCHECK3 alignment: SDIL(trained A) vs DFA(fixed A):")
+    # ---- CHECK 3: report DFA as a sanity comparator. On a single repeatedly
+    # trained batch, ordinary feedback alignment can exceed learned A, so no
+    # universal ordering is asserted here; CHECK2 tests that A actually learns.
+    print("\nCHECK3 alignment sanity comparison: SDIL(trained A) vs DFA(fixed A):")
     net_dfa = SDILNet(sizes, device=dev, seed=3)
     cfg_dfa = dfa_config(eta=0.1, momentum=0.9)
     for step in range(401):
diff --git a/sdil/core.py b/sdil/core.py
index 6137171..ce42650 100644
--- a/sdil/core.py
+++ b/sdil/core.py
@@ -214,8 +214,8 @@ def node_perturbation_targets(net, x, y, sigma=1e-2, antithetic=True, n_dirs=1):
     (B, n_l) tensors. No autograd.
 
     We perturb h_l -> h_l + sigma*xi, propagate forward THROUGH THE REST of the
-    net (weights only), read the loss change dL, and form
-        q_l = -(dL / sigma^2) * xi        (spec convention; scale absorbed by eta_A)
+    net (weights only), read the loss change dL, and form the antithetic estimate
+        q_l = -((L(+sigma*xi) - L(-sigma*xi)) / (2*sigma)) * xi
     Antithetic pairs (+xi, -xi) cancel the O(sigma^2) even term. Averaging over
     n_dirs independent directions cuts the estimator variance ~1/n_dirs, which
     the predecessor project's failures flagged as the main risk of learning a
@@ -287,9 +287,9 @@ class SDILConfig:
         self.feedback = feedback              # "error" | "error_deriv"
         self.p_update_on_neutral = p_update_on_neutral
         # node perturbation estimates the descent DIRECTION well but its
-        # magnitude (~ -dL/sigma^2) is not ||grad||, so A learns a mis-scaled
-        # per-layer step. Normalising each layer's delta to unit RMS decouples
-        # step size (set by eta) from that miscalibration -- like normalized SGD.
+        # magnitude is noisy, so A can learn a mis-scaled per-layer step.
+        # Normalising each layer's delta to unit RMS decouples step size (set by
+        # eta) from that noise -- like normalized SGD.
         self.normalize_delta = normalize_delta
 
 
@@ -327,6 +327,13 @@ def sdil_step(net, x, y, y_onehot, cfg, step, prev_error=None):
             delta = r_list[l] * gain                   # (B, n_l)
             if cfg.normalize_delta:
                 delta = delta / (delta.pow(2).mean().sqrt() + 1e-8)
+            # For an interior residual block h'=h+alpha*phi(Wh), the local
+            # Jacobian dh'/dW contains alpha.  Omitting it preserves direction
+            # but makes the effective step grow as 1/alpha=sqrt(depth), which
+            # confounds depth-scaling comparisons.  The input projection (l=0)
+            # is plain rather than residual and therefore has scale 1.
+            if net.residual and l >= 1:
+                delta = net.res_alpha * delta
             dW = delta.t() @ h[l] / B                  # (n_l, n_{l-1})
             db = delta.mean(0)
             _apply(net, l, dW, db, cfg)