summaryrefslogtreecommitdiff
path: root/notes/08_sgd_tangent_kernel_validation.md
blob: a2b193acc3054416387f59e7c85c2b331fb821af (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
# Short-SGD Tangent-Kernel Validation

This note records the first replacement for the rejected Adam downstream plot.
The goal is to validate a quantitative theory object, not to fit a post-hoc
transition curve.

## Setup

Task:

- random-label regression;
- input dimension \(16\);
- output dimension \(4\);
- \(N=128\), so \(Nn_L=512\);
- architecture \(16\to n\to n\to 4\);
- widths \(n\in\{8,12,16,24,32,48,64,96\}\).

Training:

- optimizer: SGD;
- learning rate: \(10^{-3}\);
- steps: \(50\);
- init seeds: \(4\);
- feedback seeds per init: \(8\);
- total FA trajectories: \(8\times4\times8=256\).

The short horizon is intentional. It keeps the experiment close to the fixed
tangent-kernel regime where the theory is actually supposed to apply.

## Theory Prediction

For each initialization and feedback seed, compute the full training-set BP
Jacobian

\[
J_{\mathrm{BP}}\in\mathbb R^{512\times P},
\]

and the FA pseudo-Jacobian

\[
\tilde J_{\mathrm{FA}}\in\mathbb R^{512\times P}.
\]

Then form

\[
K_{\mathrm{BP}}=J_{\mathrm{BP}}J_{\mathrm{BP}}^\top,
\qquad
K_{\mathrm{FA}}=J_{\mathrm{BP}}\tilde J_{\mathrm{FA}}^\top.
\]

For the initial residual \(r_0=f_{\theta_0}(X)-y\), predict finite-time loss by

\[
r_T^{\mathrm{lin}}
=
\exp\left(-\frac{\eta T}{N}K\right)r_0,
\]

\[
L_T^{\mathrm{lin}}
=
\frac{1}{2N}\|r_T^{\mathrm{lin}}\|^2.
\]

The predicted FA/BP train gap is

\[
\Delta L_T^{\mathrm{pred}}
=
L_T^{\mathrm{lin,FA}}-L_T^{\mathrm{lin,BP}}.
\]

No fitted scale or shifted task dimension is used.

## Results

Output directory:

`outputs/fa_tangent_kernel_capacity_sgd_short_256`

Main figures:

- `outputs/fa_tangent_kernel_capacity_sgd_short_256/sgd_short_theory_distribution_empirical_trajectories.png`
- `outputs/fa_tangent_kernel_capacity_sgd_short_256/sgd_short_theory_empirical_mean_by_margin.png`
- `outputs/fa_tangent_kernel_capacity_sgd_short_256/predicted_vs_empirical_train_gap.png`

Aggregate over 256 FA trajectories:

\[
\operatorname{corr}
\left(
\Delta L_T^{\mathrm{pred}},
\Delta L_T^{\mathrm{emp}}
\right)
=
0.9725.
\]

\[
\operatorname{MAE}
=
0.0560,
\qquad
\operatorname{RMSE}
=
0.1066.
\]

The theory overpredicts the magnitude slightly:

\[
\mathbb E[\Delta L_T^{\mathrm{pred}}]=0.2211,
\qquad
\mathbb E[\Delta L_T^{\mathrm{emp}}]=0.1651.
\]

The empirical-vs-predicted linear fit is approximately

\[
\Delta L_T^{\mathrm{emp}}
\approx
0.546\,\Delta L_T^{\mathrm{pred}}+0.044.
\]

This is a diagnostic bias, not a fitted correction for the theory.

## Interpretation

The short-SGD experiment validates the tangent-kernel prediction as a
quantitative local theory: the predicted distribution and the empirical
trajectory distribution have the same structure, and the per-trajectory
correlation is high.

The earlier 1000-step SGD probe does not match the fixed initialization kernel:
the correlation becomes negative and the prediction overestimates gap magnitude.
That failure is expected because long FA training changes the forward weights
and feedback alignment, so \(K_{\mathrm{FA}}\) is time-varying.

Therefore the experimental program should be split cleanly:

1. short/lazy SGD: fixed-kernel quantitative distribution validation;
2. long-horizon memorization transition: time-varying FA kernel or explicit
   alignment-dynamics theory.

Do not use Adam for the main quantitative validation.