fix: match published baseline protocols

3 files changed, 58 insertions, 15 deletions

diff --git a/sdil/data.py b/sdil/data.py
index f5b4928..8253d84 100644
--- a/sdil/data.py
+++ b/sdil/data.py
@@ -85,7 +85,8 @@ def _load_cifar(data_dir, n_classes=10):
             torch.from_numpy(norm(xte)), torch.from_numpy(yte))
 
 
-def get_dataset(name="mnist", batch_size=128, data_dir=DATA_DIR, device="cpu", preload=True):
+def get_dataset(name="mnist", batch_size=128, data_dir=DATA_DIR, device="cpu", preload=True,
+                shuffle_train=True, train_limit=None):
     if name == "cifar10":
         xtr, ytr, xte, yte = _load_cifar(data_dir)
         n_in = 3072
@@ -94,9 +95,13 @@ def get_dataset(name="mnist", batch_size=128, data_dir=DATA_DIR, device="cpu", p
         xtr, ytr = _load_split(subdir, True, mean, std, data_dir)
         xte, yte = _load_split(subdir, False, mean, std, data_dir)
         n_in = 784
+    if train_limit is not None:
+        if train_limit <= 0 or train_limit > xtr.shape[0]:
+            raise ValueError(f"train_limit must be in [1, {xtr.shape[0]}], got {train_limit}")
+        xtr, ytr = xtr[:train_limit], ytr[:train_limit]
     xtr, ytr = xtr.to(device), ytr.to(device)
     xte, yte = xte.to(device), yte.to(device)
-    return (_FastLoader(xtr, ytr, batch_size, True),
+    return (_FastLoader(xtr, ytr, batch_size, shuffle_train),
             _FastLoader(xte, yte, 1000, False), n_in, 10)
 
 
diff --git a/sdil/local_baselines.py b/sdil/local_baselines.py
index 23b1b5c..dc33fea 100644
--- a/sdil/local_baselines.py
+++ b/sdil/local_baselines.py
@@ -41,26 +41,41 @@ class FANet(SDILNet):
             self.B.append((torch.randn(*shape, generator=g) * (b_scale / math.sqrt(shape[1]))
                            ).to(self.W[l].device, self.dtype))
 
+    def fa_hidden_gradients(self, fwd, yoh):
+        """Return FA estimates of dL/dh for each hidden block output.
+
+        The returned tensors include the residual-branch scale used by the
+        corresponding local weight update.  Identity skips transport ``dh``
+        exactly; only the nonlinear branch is transported through fixed B.
+        Keeping this calculation separate makes the signal directly
+        measurable without using autograd for learning.
+        """
+        h, u = fwd["h"], fwd["u"]
+        e = torch.softmax(h[-1], 1) - yoh
+        dh = e @ self.B[self.L - 1]
+        estimates = [None] * (self.L - 1)
+        for l in range(self.L - 2, -1, -1):
+            branch_scale = self.res_alpha if self.residual and l >= 1 else 1.0
+            estimates[l] = branch_scale * dh
+            du = estimates[l] * self.act_prime(u[l])
+            if l > 0:
+                feedback = du @ self.B[l]
+                dh = dh + feedback if self.residual and l >= 1 else feedback
+        return estimates
+
     def fa_step(self, x, y, yoh, eta, momentum=0.0):
         with torch.no_grad():
             fwd = self.forward(x)
             h, u = fwd["h"], fwd["u"]
             e = torch.softmax(h[-1], 1) - yoh          # grad wrt logits
             loss = F.cross_entropy(h[-1], y).item()
+            hidden_grad = self.fa_hidden_gradients(fwd, yoh)
             # output layer (exact local)
             self._upd(self.L - 1, e, h[self.L - 1], eta, momentum)
 
-            # dh is the error wrt the hidden *state*.  On a residual block,
-            # h'=h+alpha*phi(Wh), the identity branch transports dh exactly and
-            # only the nonlinear branch uses a fixed random feedback matrix.
-            dh = e @ self.B[self.L - 1]
-            for l in range(self.L - 2, -1, -1):
-                branch_scale = self.res_alpha if self.residual and l >= 1 else 1.0
-                du = branch_scale * dh * self.act_prime(u[l])
+            for l in range(self.L - 1):
+                du = hidden_grad[l] * self.act_prime(u[l])
                 self._upd(l, du, h[l], eta, momentum)
-                if l > 0:
-                    feedback = du @ self.B[l]
-                    dh = dh + feedback if self.residual and l >= 1 else feedback
             return loss
 
     def _upd(self, l, delta, pre, eta, momentum):
@@ -321,9 +336,13 @@ class EPNet:
             s = new
         return s
 
-    def train_step(self, x, y, yoh, eta):
+    def train_step(self, x, y, yoh, eta, free_state=None, return_free_state=False):
         with torch.no_grad():
-            s0 = self._settle(x, beta=0.0, T=self.T_free)                 # free phase
+            s0 = self._settle(x, beta=0.0, s=free_state, T=self.T_free)   # free phase
+            # The authors keep one free-phase particle per training example
+            # and reuse it on the next epoch.  Save the free state, not the
+            # subsequently nudged state.
+            persistent_state = [t.detach().clone() for t in s0]
             beta = self.beta
             if self.random_beta_sign and torch.randint(0, 2, ()).item() == 0:
                 beta = -beta
@@ -338,7 +357,8 @@ class EPNet:
                 self.W[k] += layer_eta * dW
                 self.b[k] += layer_eta * db
             # free-phase output as prediction proxy for loss logging
-            return F.mse_loss(s0[-1], yoh).item()
+            loss = F.mse_loss(s0[-1], yoh).item()
+            return (loss, persistent_state) if return_free_state else loss
 
     @torch.no_grad()
     def predict(self, x):
diff --git a/sdil/probes.py b/sdil/probes.py
index fad82ca..2d10aad 100644
--- a/sdil/probes.py
+++ b/sdil/probes.py
@@ -85,6 +85,24 @@ def alignment_report(net, x, y, y_onehot, cfg):
 
 
 @torch.no_grad()
+def fa_alignment_report(net, x, y, y_onehot):
+    """Alignment of sequential-FA local update signals with exact descent.
+
+    ``FANet.fa_hidden_gradients`` performs the same fixed-feedback transport
+    used during learning.  Autograd appears only here to obtain the reference
+    hidden-state gradients for measurement.
+    """
+    grads, loss = true_hidden_grads(net, x, y)
+    fwd = net.forward(x)
+    estimated = net.fa_hidden_gradients(fwd, y_onehot)
+    cosines = []
+    for l, (estimate, true_grad) in enumerate(zip(estimated, grads)):
+        gain = net.act_prime(fwd["u"][l])
+        cosines.append(_row_cos(-estimate * gain, -true_grad * gain))
+    return {"cos_fa_negg": cosines, "loss": loss}
+
+
+@torch.no_grad()
 def loss_decrease_ratio(net, x, y, y_onehot, cfg, step):
     """Single-step descent quality: apply one SDIL update to a scratch copy and
     measure the actual loss drop on the same batch, compared to one plain-SGD