Derive actual FA initial operator moment

1 files changed, 327 insertions, 0 deletions

diff --git a/notes/29_actual_fa_initial_operator_moment_theorem.md b/notes/29_actual_fa_initial_operator_moment_theorem.md
new file mode 100644
index 0000000..5814f06
--- /dev/null
+++ b/notes/29_actual_fa_initial_operator_moment_theorem.md
@@ -0,0 +1,327 @@
+# Actual FA Initial Operator Moment Theorem
+
+This note records the first successful attempt to replace the hard-k
+random-subspace proxy with a theorem about actual FA.
+
+## Problem
+
+The hard random-subspace null model says:
+
+```text
+E[e] = k/P
+```
+
+but experiments showed it overpredicts real FA gaps. The reason is that actual
+FA does not delete a hard `k`-dimensional subspace. It changes the backward
+operator.
+
+We need a quantity closer to real FA.
+
+## Setup
+
+Squared loss with residual vector `r`.
+
+For each parameter block/layer `l`:
+
+```text
+g_l^BP = BP gradient block
+g_l^FA = FA pseudo-gradient block
+```
+
+The BP one-step learning speed is:
+
+```text
+speed_BP = sum_l ||g_l^BP||^2
+```
+
+The FA first-order loss decrease uses the mixed inner product:
+
+```text
+speed_FA = sum_l <g_l^BP, g_l^FA>
+```
+
+Equivalently:
+
+```text
+speed_BP = r^T K_BP r
+speed_FA = r^T K_FA r
+```
+
+where:
+
+```text
+K_BP = J J^T
+K_FA = J J_tilde^T
+```
+
+The local operator erosion is:
+
+```text
+e_0(B) = 1 - speed_FA / speed_BP
+```
+
+## Theorem: Exact Conditional Mean at Initialization
+
+Assume:
+
+```text
+1. forward weights W are fixed;
+2. residual r is fixed;
+3. feedback matrices B are independent of W and r;
+4. each feedback matrix has zero-mean entries or, more generally, zero
+   first moment in every hidden backward direction;
+5. the output layer is trained with the true output gradient, as in standard FA.
+```
+
+Then, at initialization:
+
+```text
+E_B[speed_FA | W, r] = ||g_output^BP||^2
+```
+
+Therefore:
+
+```text
+E_B[ speed_FA / speed_BP | W, r ]
+  = ||g_output^BP||^2 / sum_l ||g_l^BP||^2
+```
+
+and:
+
+```text
+E_B[e_0 | W, r]
+  = 1 - ||g_output^BP||^2 / sum_l ||g_l^BP||^2
+```
+
+So the initial effective FA burden is not raw hard constraint count. It is the
+fraction of BP learning speed carried by hidden-layer gradient blocks:
+
+```text
+e_0_mean = hidden BP speed share
+```
+
+If we want to express it as an effective dimension:
+
+```text
+k_eff,0 = P * e_0_mean
+```
+
+but the more natural object is the dimensionless erosion `e_0_mean`.
+
+## Proof Sketch
+
+The output-layer FA gradient equals the BP gradient:
+
+```text
+g_output^FA = g_output^BP
+```
+
+so its mixed contribution is:
+
+```text
+<g_output^BP, g_output^FA> = ||g_output^BP||^2
+```
+
+Every hidden FA pseudo-gradient contains at least one zero-mean independent
+feedback matrix. Conditional on fixed forward weights and residuals:
+
+```text
+E_B[g_l^FA | W, r] = 0
+```
+
+for hidden layers `l`.
+
+Thus:
+
+```text
+E_B[<g_l^BP, g_l^FA> | W, r] = 0
+```
+
+for every hidden layer.
+
+Adding the layer contributions gives the result.
+
+## Relation to Song, Xu, and Lafferty
+
+For a two-layer scalar-output network, Song/Xu/Lafferty decompose FA dynamics
+as:
+
+```text
+K_FA = G + H_FA
+```
+
+where `G` is the top-layer kernel and `H_FA` is the feedback-dependent hidden
+term. This theorem is the residual-direction version of that decomposition at
+initialization:
+
+```text
+E_B[H_FA | W, r] = 0
+```
+
+so:
+
+```text
+E_B[K_FA | W] = K_output
+```
+
+not:
+
+```text
+E_B[K_FA | W] = (1 - k/P) K_BP
+```
+
+This explains why the hard-k random-subspace theory was too pessimistic.
+
+## Distribution, Not Just Mean
+
+For one-hidden-layer scalar-output FA with Gaussian feedback, the hidden mixed
+speed is exactly Gaussian conditional on `W` and `r`.
+
+It can be written as:
+
+```text
+speed_FA = output_speed + Z
+```
+
+where:
+
+```text
+E_B[Z | W, r] = 0
+Var_B[Z | W, r] = feedback_variance * ||C(W,r)||^2
+```
+
+for an explicit coefficient tensor `C(W,r)` determined by activations, gates,
+residuals, and BP hidden gradients.
+
+Therefore:
+
+```text
+e_0(B) = 1 - (output_speed + Z) / speed_BP
+```
+
+has an explicit conditional Gaussian distribution in the one-hidden-layer
+Gaussian-feedback case.
+
+For deeper FA, hidden terms contain products of feedback matrices. The exact
+mean above still holds, but the finite-width distribution is no longer exactly
+Gaussian. A CLT/Wick expansion should give a computable approximation.
+
+## Infinite-Width Interpretation
+
+At infinite width, layerwise BP speed shares should converge to deterministic
+NTK/NNGP quantities:
+
+```text
+output_share = r^T K_output r / r^T K_BP r
+hidden_share = 1 - output_share
+```
+
+Thus an architecture-level prediction may be obtained from layerwise NTK
+recursions:
+
+```text
+E_B[e_0] -> hidden_share_NTK
+```
+
+This gives a principled route to an architecture/data-level prediction without
+post-hoc scaling.
+
+## Time-Varying Extension
+
+At later time:
+
+```text
+speed_FA,t = output_speed_t + hidden_alignment_gain_t
+```
+
+where:
+
+```text
+hidden_alignment_gain_t = sum_hidden <g_l^BP(t), g_l^FA(t)>
+```
+
+At initialization:
+
+```text
+E_B[hidden_alignment_gain_0 | W,r] = 0
+```
+
+During training, forward weights become dependent on feedback matrices, so this
+term can become positive. This is actual feedback alignment.
+
+Therefore:
+
+```text
+e_t = 1 - (output_speed_t + hidden_alignment_gain_t) / speed_BP,t
+```
+
+The current tangent-operator estimator works because it measures early movement
+of this alignment gain. A future theorem should bound or approximate:
+
+```text
+hidden_alignment_gain_t
+```
+
+from the early dynamics, possibly using the perturbative `H_FA` analysis style
+from Song/Xu/Lafferty.
+
+## Empirical Validation
+
+Script:
+
+```text
+scripts/actual_fa_initial_operator_moments.py
+```
+
+Run:
+
+```text
+python scripts/actual_fa_initial_operator_moments.py \
+  --widths 16 24 32 48 64 96 \
+  --init-seeds 6 \
+  --feedback-samples 512 \
+  --torch-threads 8 \
+  --outdir outputs/actual_fa_initial_operator_moments
+```
+
+Outputs:
+
+```text
+outputs/actual_fa_initial_operator_moments/initial_operator_moment_rows.csv
+outputs/actual_fa_initial_operator_moments/predicted_vs_empirical_initial_erosion_mean.png
+outputs/actual_fa_initial_operator_moments/initial_erosion_mean_error_by_width.png
+outputs/actual_fa_initial_operator_moments/initial_erosion_histograms.png
+```
+
+Representative results:
+
+| width | init | theory mean erosion | empirical mean erosion | error |
+|---:|---:|---:|---:|---:|
+| 16 | 0 | 0.451931 | 0.444348 | -0.00758 |
+| 24 | 0 | 0.312016 | 0.312286 | 0.00027 |
+| 32 | 0 | 0.236058 | 0.239714 | 0.00366 |
+| 48 | 0 | 0.261471 | 0.259858 | -0.00161 |
+| 64 | 0 | 0.209125 | 0.211261 | 0.00214 |
+| 96 | 0 | 0.259765 | 0.259147 | -0.00062 |
+
+Across all tested widths and initializations, the no-fit mean prediction is
+accurate to Monte Carlo error.
+
+## Consequence for Contributions
+
+This improves contribution 2:
+
+```text
+hard-k random-subspace erosion = clean null model
+actual FA initial erosion = hidden BP speed share theorem
+```
+
+Contribution 4 still handles finite-time gap:
+
+```text
+initial theorem gives e_0
+operator estimator tracks e_t and converts it into finite-time gap
+```
+
+This is a stronger and cleaner story than trying to force raw capacity margin
+to predict the full trajectory gap.