Switch e0 distribution sampler to real FA backward

The sampler previously drew e0 through the analytic linear form <C,B>,
making the Gaussian KS test near-tautological. It now runs the real FA
backward pass per feedback draw and reports the max deviation from the
linear identity (<= 8.9e-16). Rerun reproduces the published table.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

1 files changed, 45 insertions, 9 deletions

diff --git a/scripts/actual_fa_initial_erosion_distribution.py b/scripts/actual_fa_initial_erosion_distribution.py
index 019ec15..7f469a1 100644
--- a/scripts/actual_fa_initial_erosion_distribution.py
+++ b/scripts/actual_fa_initial_erosion_distribution.py
@@ -10,6 +10,13 @@ with Gaussian feedback B. Therefore the one-step erosion
     e0 = 1 - speed_FA / speed_BP
 
 is exactly Gaussian conditional on the forward weights and data.
+
+Sampling runs the *real FA backward pass* for every feedback draw
+(delta_hidden = (delta_out @ B^T) * gate, then <g_hidden_BP, g_hidden_FA>),
+so the empirical histogram is independent of the analytic linear form
+<C, B>. The linear identity is checked against the real backward per draw
+and its max absolute deviation is reported (`linear_identity_max_err`),
+instead of being used to generate the samples.
 """
 
 from __future__ import annotations
@@ -45,6 +52,7 @@ class DistributionRow:
     ks_pvalue: float
     mean_error: float
     std_error: float
+    linear_identity_max_err: float
 
 
 def parse_args() -> argparse.Namespace:
@@ -136,26 +144,50 @@ def theory_terms(
 
 
 def sample_erosions(
-    bp_speed: float,
-    output_speed: float,
+    w1: Tensor,
+    w2: Tensor,
+    x: Tensor,
+    y: Tensor,
     coeff: Tensor,
     scale: float,
     samples: int,
     seed: int,
-) -> np.ndarray:
+) -> tuple[np.ndarray, float]:
+    """Sample e0 by running the real FA backward pass per feedback draw.
+
+    Returns the erosion samples and the max abs deviation between the real
+    mixed term <g_hidden_BP, g_hidden_FA(B)> and the analytic linear form
+    <C, B> (an internal consistency check; it is NOT used for sampling).
+    """
+    z, h, out = forward(w1, w2, x)
+    batch = x.shape[0]
+    delta_out = (out - y) / batch
+    gate = (z > 0).to(torch.float64)
+    grad_out = delta_out.T @ h
+    delta_hidden_bp = (delta_out @ w2) * gate
+    grad_hidden_bp = delta_hidden_bp.T @ x
+    output_speed = float(torch.sum(grad_out * grad_out))
+    bp_speed = output_speed + float(torch.sum(grad_hidden_bp * grad_hidden_bp))
+
     generator = torch.Generator(device="cpu")
     generator.manual_seed(seed)
     values = []
+    identity_err = 0.0
     for _ in range(samples):
         feedback = torch.randn(
             coeff.shape,
             generator=generator,
             dtype=torch.float64,
         ) * scale
-        hidden_mixed = float(torch.sum(feedback * coeff))
+        # real FA backward: B replaces W2^T in the hidden delta
+        delta_hidden_fa = (delta_out @ feedback.T) * gate
+        grad_hidden_fa = delta_hidden_fa.T @ x
+        hidden_mixed = float(torch.sum(grad_hidden_bp * grad_hidden_fa))
+        hidden_mixed_linear = float(torch.sum(feedback * coeff))
+        identity_err = max(identity_err, abs(hidden_mixed - hidden_mixed_linear))
         speed_fa = output_speed + hidden_mixed
         values.append(1.0 - speed_fa / bp_speed)
-    return np.array(values, dtype=np.float64)
+    return np.array(values, dtype=np.float64), identity_err
 
 
 def run_one(args: argparse.Namespace, x: Tensor, y: Tensor, width: int, init_seed: int) -> tuple[DistributionRow, np.ndarray]:
@@ -168,9 +200,11 @@ def run_one(args: argparse.Namespace, x: Tensor, y: Tensor, width: int, init_see
         y,
         scale,
     )
-    erosions = sample_erosions(
-        bp_speed,
-        output_speed,
+    erosions, identity_err = sample_erosions(
+        w1,
+        w2,
+        x,
+        y,
         coeff,
         scale,
         args.feedback_samples,
@@ -193,6 +227,7 @@ def run_one(args: argparse.Namespace, x: Tensor, y: Tensor, width: int, init_see
         ks_pvalue=float(ks.pvalue),
         mean_error=float(np.mean(erosions) - theory_mean),
         std_error=float(np.std(erosions, ddof=1) - theory_std),
+        linear_identity_max_err=identity_err,
     )
     return row, erosions
 
@@ -305,7 +340,8 @@ def main() -> None:
                 f"width={width} init={init_seed}: "
                 f"mean theory={row.theory_mean:.6f} empirical={row.empirical_mean:.6f}; "
                 f"std theory={row.theory_std:.6f} empirical={row.empirical_std:.6f}; "
-                f"KS={row.ks_stat:.4f} p={row.ks_pvalue:.3g}",
+                f"KS={row.ks_stat:.4f} p={row.ks_pvalue:.3g}; "
+                f"identity_err={row.linear_identity_max_err:.3g}",
                 flush=True,
             )