Investigate finite-T kernel errors

1 files changed, 363 insertions, 0 deletions

diff --git a/notes/12_finite_t_investigation.md b/notes/12_finite_t_investigation.md
new file mode 100644
index 0000000..0465129
--- /dev/null
+++ b/notes/12_finite_t_investigation.md
@@ -0,0 +1,363 @@
+# Finite-T Investigation: Why BP Is Estimated Too Strong And FA Too Weak
+
+This note investigates the finite-time regime beyond the tagged local result
+`t5-local-tangent-validation`.
+
+## Empirical Decomposition
+
+For larger \(T\), decompose the gap error into BP and FA terms:
+
+\[
+\Delta_{\mathrm{gap}}^{\mathrm{err}}
+=
+\left(L_{\mathrm{FA}}^{\mathrm{pred}}-L_{\mathrm{BP}}^{\mathrm{pred}}\right)
+-
+\left(L_{\mathrm{FA}}^{\mathrm{emp}}-L_{\mathrm{BP}}^{\mathrm{emp}}\right).
+\]
+
+Equivalently,
+
+\[
+\Delta_{\mathrm{gap}}^{\mathrm{err}}
+=
+\underbrace{
+\left(L_{\mathrm{FA}}^{\mathrm{pred}}-L_{\mathrm{FA}}^{\mathrm{emp}}\right)
+}_{\mathrm{FA\ error}}
+-
+\underbrace{
+\left(L_{\mathrm{BP}}^{\mathrm{pred}}-L_{\mathrm{BP}}^{\mathrm{emp}}\right)
+}_{\mathrm{BP\ error}}.
+\]
+
+Observed signs:
+
+- fixed-width \(T=50\): BP is estimated too strong; FA is estimated slightly too
+  weak;
+- width-sweep \(T=1000\): BP is estimated too strong; FA is estimated much too
+  weak.
+
+Numerically:
+
+| experiment | BP pred-emp | FA pred-emp | gap pred-emp |
+|---|---:|---:|---:|
+| fixed-width \(N=128,T=50\) | -0.0233 | +0.0052 | +0.0285 |
+| fixed-width \(N=320,T=50\) | -0.0167 | +0.0050 | +0.0217 |
+| width-sweep \(T=1000\) | -0.1471 | +0.3652 | +0.5123 |
+
+So:
+
+1. at moderate \(T\), the dominant error is BP fixed-kernel optimism;
+2. at long \(T\), FA fixed-kernel pessimism becomes dominant because FA learns
+   alignment.
+
+## Related Work
+
+The original NTK theory shows that in the infinite-width limit, gradient
+descent is described by a kernel gradient flow and the NTK stays constant during
+training. For finite networks the NTK depends on parameters and can change
+during training.
+
+Chizat, Oyallon, and Bach's lazy-training view makes the same distinction:
+linearized/lazy dynamics are accurate when parameters barely move; outside that
+regime, feature learning changes the tangent kernel.
+
+Finite-width studies ask when NTK theory can be trusted. Seleznova and Kutyniok
+emphasize that the empirical NTK can change substantially in finite deep
+networks, depending on initialization and depth.
+
+The Neural Tangent Hierarchy of Huang and Yau is the closest formal tool for
+finite-width kernel evolution: it describes the time derivative of the NTK using
+higher-order tangent objects. This is exactly the type of correction needed for
+finite \(T\).
+
+For FA/DFA specifically, Refinetti et al.'s "align, then memorise" picture says
+that feedback-alignment learning first adapts weights to improve gradient
+alignment, then fits data. This explains why fixed \(K_{\mathrm{FA}}(0)\) becomes
+too pessimistic at long \(T\).
+
+Surrogate-gradient NTK work is relevant because FA is a surrogate-gradient
+learning rule: the forward Jacobian controls output motion while the surrogate
+backward Jacobian controls parameter motion.
+
+## Own Derivation
+
+Use squared loss
+
+\[
+L(\theta)
+=
+\frac{1}{2N}\|r(\theta)\|^2,
+\qquad
+r(\theta)=f_\theta(X)-y.
+\]
+
+Let
+
+\[
+J_t
+=
+\frac{\partial f_\theta(X)}{\partial\theta}
+\bigg|_{\theta_t}
+\]
+
+be the true forward Jacobian, and let
+
+\[
+\tilde J_t
+\]
+
+be the surrogate Jacobian used by the learning rule. For BP,
+
+\[
+\tilde J_t=J_t.
+\]
+
+For FA,
+
+\[
+\tilde J_t=\tilde J_{\mathrm{FA}}(B,\theta_t).
+\]
+
+Full-batch update:
+
+\[
+\theta_{t+1}
+=
+\theta_t
+-
+\frac{\eta}{N}
+\tilde J_t^\top r_t.
+\]
+
+Taylor expand outputs:
+
+\[
+r_{t+1}
+=
+r_t
++
+J_t(\theta_{t+1}-\theta_t)
++
+\frac12
+H_{f,t}[
+\theta_{t+1}-\theta_t,
+\theta_{t+1}-\theta_t
+]
++
+O(\|\Delta\theta_t\|^3).
+\]
+
+Substitute the update:
+
+\[
+r_{t+1}
+=
+\left(
+I-\frac{\eta}{N}K_t
+\right)r_t
++
+\frac{\eta^2}{2N^2}
+H_{f,t}
+[
+\tilde J_t^\top r_t,
+\tilde J_t^\top r_t
+]
++
+O(\eta^3),
+\]
+
+where the learning-rule kernel is
+
+\[
+K_t=J_t\tilde J_t^\top.
+\]
+
+The local \(T=5\) theorem uses
+
+\[
+K_t\approx K_0
+\]
+
+and ignores the \(O(\eta^2)\) Taylor term:
+
+\[
+r_T^{(0)}
+=
+\prod_{t=0}^{T-1}
+\left(I-\frac{\eta}{N}K_0\right)r_0.
+\]
+
+The finite-\(T\) correction has two pieces:
+
+\[
+\delta r_T
+\approx
+\sum_{s<T}
+A_{T:s+1}
+\left[
+-\frac{\eta}{N}\Delta K_s r_s^{(0)}
++
+\frac{\eta^2}{2N^2}
+H_{f,s}[
+\tilde J_s^\top r_s^{(0)},\tilde J_s^\top r_s^{(0)}
+]
+\right],
+\]
+
+where
+
+\[
+A_{T:s+1}
+=
+\prod_{u=s+1}^{T-1}
+\left(I-\frac{\eta}{N}K_0\right),
+\qquad
+\Delta K_s=K_s-K_0.
+\]
+
+So finite-\(T\) error is not mysterious: it is accumulated kernel drift plus
+second-order output curvature.
+
+## Why BP Is Estimated Too Strong
+
+For BP:
+
+\[
+K_t^{\mathrm{BP}}=J_tJ_t^\top.
+\]
+
+The fixed-kernel prediction assumes the initial tangent features keep their
+initial training power. In the experiments, \(L_{\mathrm{BP}}^{\mathrm{pred}}\)
+is below \(L_{\mathrm{BP}}^{\mathrm{emp}}\), so the true finite-width BP dynamics
+learn more slowly than the frozen \(K_0\) model predicts.
+
+A sufficient explanation is that along the residual direction,
+
+\[
+r_t^\top K_t^{\mathrm{BP}} r_t
+<
+r_t^\top K_0^{\mathrm{BP}} r_t
+\]
+
+on average, or that curvature terms partially oppose the first-order output
+motion. This is finite-width/lazy-regime breakdown: the initial BP NTK is too
+optimistic as a long-horizon proxy.
+
+This matches the literature: finite-width networks can leave the frozen-NTK
+regime, and the empirical NTK can change significantly during training.
+
+## Why FA Is Estimated Too Weak
+
+For FA:
+
+\[
+K_t^{\mathrm{FA}}
+=
+J_t\tilde J_t^\top.
+\]
+
+At \(t=0\), \(\tilde J_t\) is based on random feedback and is misaligned with
+the BP gradient. The fixed-kernel prediction freezes this bad credit-assignment
+operator.
+
+But FA training changes the forward weights so that the surrogate direction
+becomes more useful. In kernel language, the alignment improves the real part
+or residual-direction action of \(K_t^{\mathrm{FA}}\):
+
+\[
+r_t^\top K_t^{\mathrm{FA}} r_t
+>
+r_t^\top K_0^{\mathrm{FA}} r_t
+\]
+
+on average after an alignment phase.
+
+Thus \(L_{\mathrm{FA}}^{\mathrm{pred}}>L_{\mathrm{FA}}^{\mathrm{emp}}\): the
+frozen random-feedback kernel underestimates FA learning.
+
+This is exactly consistent with "align, then memorise": early training improves
+the credit-assignment geometry before memorisation dominates.
+
+## Practical Finite-T Theory Candidate
+
+The next theory should not be a fitted multiplicative scale. The clean finite-T
+object is:
+
+\[
+r_T^{\mathrm{tv}}
+=
+\prod_{t=0}^{T-1}
+\left(
+I-\frac{\eta}{N}K_t
+\right)r_0.
+\]
+
+This is directly testable by measuring \(K_t\) at checkpoints.
+
+If the product using measured \(K_t\) matches empirical \(T=50\) and \(T=1000\),
+then the remaining task is to model \(K_t\). A compressed model could be:
+
+\[
+K_t^{\mathrm{FA}}
+\approx
+K_0^{\mathrm{FA}}
++
+\alpha_t
+\left(
+K_t^{\mathrm{BP}}-K_0^{\mathrm{FA}}
+\right)
++
+E_t,
+\]
+
+where \(\alpha_t\) is not fitted freely but defined from measurable alignment:
+
+\[
+\alpha_t
+=
+\frac{
+\langle K_t^{\mathrm{FA}}-K_0^{\mathrm{FA}},
+K_t^{\mathrm{BP}}-K_0^{\mathrm{FA}}\rangle_F
+}{
+\|K_t^{\mathrm{BP}}-K_0^{\mathrm{FA}}\|_F^2
+}.
+\]
+
+Then one can test whether \(\alpha_t\) itself follows a simple law, for example
+an exponential alignment curve:
+
+\[
+\alpha_t\approx 1-\exp(-t/\tau_{\mathrm{align}}).
+\]
+
+But this is only acceptable if \(\tau_{\mathrm{align}}\) can be predicted from
+architecture/data/initialization statistics or independently measured, not
+chosen to fit final loss.
+
+## Recommended Diagnostic
+
+For a small set of trajectories:
+
+1. save checkpoints at \(t=\{0,1,2,5,10,20,50,100,\dots\}\);
+2. compute \(K_t^{\mathrm{BP}}\), \(K_t^{\mathrm{FA}}\), residuals \(r_t\);
+3. measure directional learning rates
+
+\[
+\lambda_t^{\mathrm{dir}}
+=
+\frac{r_t^\top K_t r_t}{\|r_t\|^2};
+\]
+
+4. measure kernel overlap
+
+\[
+\rho_t
+=
+\frac{\langle K_t^{\mathrm{FA}},K_t^{\mathrm{BP}}\rangle_F}
+{\|K_t^{\mathrm{FA}}\|_F\|K_t^{\mathrm{BP}}\|_F};
+\]
+
+5. compare empirical residuals with the measured time-varying kernel product.
+
+This will tell us whether finite-\(T\) is mostly kernel drift or whether output
+curvature terms must be modeled explicitly.