Add hidden gradient alignment metric

3 files changed, 48 insertions, 0 deletions

diff --git a/notes/02_experiment_notes.md b/notes/02_experiment_notes.md
index 1d641e1..c44e0d0 100644
--- a/notes/02_experiment_notes.md
+++ b/notes/02_experiment_notes.md
@@ -331,6 +331,7 @@ Result:
 - BP final loss: `0.60596695`
 - FA final gap to BP: mean `0.2032423`, min `0.16939124`, max `0.25930484`
 - FA final BP/FA gradient cosine: mean `0.31281339`, min `0.29023733`, max `0.35150802`
+- FA final hidden-only BP/FA gradient cosine: mean `-0.0070103243`, min `-0.13335294`, max `0.072182807`
 
 Per-seed summary:
 
@@ -339,3 +340,5 @@ Per-seed summary:
 - seed `52`: final loss `0.78699777`, final gap `0.18103082`, initial \(Q\) mean `0.02377742`, final \(Q\) mean `0.00360508`
 
 This is only a smoke trajectory, not yet an ensemble result. It verifies that the logging pipeline can capture loss gaps, surrogate-gradient alignment, and weight-feedback alignment \(Q_l(t)\) from the same run.
+
+Important metric note: full-model gradient cosine can be inflated by the output layer, whose gradient is identical under BP and FA. Hidden-only gradient cosine is a sharper metric for feedback-induced mismatch.
diff --git a/scripts/README.md b/scripts/README.md
index 42633d0..2d7a34c 100644
--- a/scripts/README.md
+++ b/scripts/README.md
@@ -125,6 +125,7 @@ a synthetic regression task. At each checkpoint, the script records:
 
 - training loss;
 - full-model cosine between the BP gradient and the FA surrogate gradient at the FA weights;
+- hidden-layer-only cosine between the BP and FA gradients, excluding the output layer where gradients are identical;
 - layerwise \(Q_l=\cos^2(W_{l+1}^{\top},B_l)\).
 
 Outputs are written under `outputs/trajectory_mlp_fa/`:
diff --git a/scripts/trajectory_mlp_fa.py b/scripts/trajectory_mlp_fa.py
index ce20fa5..3b9d5b7 100755
--- a/scripts/trajectory_mlp_fa.py
+++ b/scripts/trajectory_mlp_fa.py
@@ -44,6 +44,7 @@ class TrajectoryRow:
     step: int
     loss: float
     gradient_cosine: float | None
+    hidden_gradient_cosine: float | None
     q_mean: float | None
     q_min: float | None
     q_max: float | None
@@ -67,6 +68,8 @@ class RunSummary:
     final_gap_to_bp: float
     initial_gradient_cosine: float | None
     final_gradient_cosine: float | None
+    initial_hidden_gradient_cosine: float | None
+    final_hidden_gradient_cosine: float | None
     initial_q_mean: float | None
     final_q_mean: float | None
 
@@ -309,6 +312,7 @@ def evaluate_bp(weights: list[Array], x: Array, y: Array, step: int) -> Trajecto
         step=step,
         loss=loss,
         gradient_cosine=1.0,
+        hidden_gradient_cosine=1.0,
         q_mean=None,
         q_min=None,
         q_max=None,
@@ -326,6 +330,7 @@ def evaluate_fa(
     bp_grads, loss = gradients(weights, x, y, feedback=None)
     fa_grads, _ = gradients(weights, x, y, feedback=feedback)
     grad_cos = cosine(flatten(bp_grads), flatten(fa_grads))
+    hidden_grad_cos = cosine(flatten(bp_grads[:-1]), flatten(fa_grads[:-1]))
     layer_cosines = layer_gradient_cosines(bp_grads, fa_grads)
     q_values = layer_q_alignments(weights, feedback)
 
@@ -335,6 +340,7 @@ def evaluate_fa(
         step=step,
         loss=loss,
         gradient_cosine=grad_cos,
+        hidden_gradient_cosine=hidden_grad_cos,
         q_mean=float(np.mean(q_values)),
         q_min=float(np.min(q_values)),
         q_max=float(np.max(q_values)),
@@ -426,6 +432,8 @@ def make_summaries(
             final_gap_to_bp=0.0,
             initial_gradient_cosine=1.0,
             final_gradient_cosine=1.0,
+            initial_hidden_gradient_cosine=1.0,
+            final_hidden_gradient_cosine=1.0,
             initial_q_mean=None,
             final_q_mean=None,
         )
@@ -442,6 +450,8 @@ def make_summaries(
                 final_gap_to_bp=final.loss - bp_final,
                 initial_gradient_cosine=first.gradient_cosine,
                 final_gradient_cosine=final.gradient_cosine,
+                initial_hidden_gradient_cosine=first.hidden_gradient_cosine,
+                final_hidden_gradient_cosine=final.hidden_gradient_cosine,
                 initial_q_mean=first.q_mean,
                 final_q_mean=final.q_mean,
             )
@@ -526,6 +536,29 @@ def save_plots(trajectories: list[TrajectoryRow], outdir: Path) -> list[Path]:
     plt.close()
     paths.append(gamma_path)
 
+    hidden_gamma_path = outdir / "hidden_gradient_cosine_curves.png"
+    plt.figure(figsize=(7, 4.5))
+    for seed in fa_seeds:
+        rows = [
+            row
+            for row in trajectories
+            if row.run_type == "fa" and row.feedback_seed == seed
+        ]
+        plt.plot(
+            [row.step for row in rows],
+            [row.hidden_gradient_cosine for row in rows],
+            alpha=0.75,
+            label=f"seed={seed}",
+        )
+    plt.axhline(0.0, color="black", linewidth=1)
+    plt.xlabel("step")
+    plt.ylabel("cos(BP hidden gradient, FA hidden gradient)")
+    plt.title("Hidden-layer surrogate gradient alignment")
+    plt.tight_layout()
+    plt.savefig(hidden_gamma_path, dpi=180)
+    plt.close()
+    paths.append(hidden_gamma_path)
+
     q_path = outdir / "q_alignment_curves.png"
     plt.figure(figsize=(7, 4.5))
     for seed in fa_seeds:
@@ -588,6 +621,11 @@ def main() -> None:
     fa_final_gammas = [
         row.final_gradient_cosine for row in summaries if row.run_type == "fa"
     ]
+    fa_final_hidden_gammas = [
+        row.final_hidden_gradient_cosine
+        for row in summaries
+        if row.run_type == "fa"
+    ]
     print(f"widths: {widths}")
     print(f"bp_final_loss: {bp_final:.8g}")
     print(
@@ -602,6 +640,12 @@ def main() -> None:
         f"min={np.min(fa_final_gammas):.8g}, "
         f"max={np.max(fa_final_gammas):.8g}"
     )
+    print(
+        "fa_final_hidden_gradient_cosine: "
+        f"mean={np.mean(fa_final_hidden_gammas):.8g}, "
+        f"min={np.min(fa_final_hidden_gammas):.8g}, "
+        f"max={np.max(fa_final_hidden_gammas):.8g}"
+    )
     print(f"summary: {outdir / 'summary.csv'}")
     print(f"trajectories: {outdir / 'trajectories.csv'}")
     print(f"layer_metrics: {outdir / 'layer_metrics.csv'}")