fix: match published baseline protocols

1 files changed, 18 insertions, 0 deletions

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