summaryrefslogtreecommitdiff
path: root/notes/04_downstream_capacity_design.md
blob: c2b97043f96b4c83e7f9f7b2808a9237b54a2a47 (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
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
# Downstream Capacity Experiment Design

## Goal

Validate the core downstream claim:

> FA and BP should behave similarly while redundant parameter capacity can absorb the FA alignment burden, and FA should degrade faster once redundancy is exhausted.

For random-label memorization, the primary performance metric is train MSE. The
experiment should produce two performance curves:

\[
L_{\mathrm{BP}}(P)
\quad \text{and} \quad
L_{\mathrm{FA}}(P),
\]

where \(P\) is model parameter count. The target signature is a widening FA/BP
train gap as \(P\) is reduced, with the widening centered near the predicted FA
capacity margin.

## Main Task Choice: Random-Label Memorization

Use:

\[
x_i\sim \mathcal N(0,I_{n_0}),
\qquad
y_i\sim \mathcal N(0,I_{n_L}).
\]

The target is train-set memorization, not test generalization. This makes the
task dimension directly measurable:

\[
d_{\mathrm{task}}\approx N n_L.
\]

For a two-hidden-layer MLP with no bias,

\[
P(n)=n_0 n+n^2+n n_L.
\]

For ordinary FA, the hard alignment-burden proxy is:

\[
K_{\mathrm{FA}}(n)
=
(n^2-1)+(n n_L-1).
\]

The predicted thresholds are:

\[
P(n)\gtrsim Nn_L
\qquad \text{for BP},
\]

and

\[
P(n)-K_{\mathrm{FA}}(n)\gtrsim Nn_L
\qquad \text{for FA}.
\]

This is the cleanest task for testing the redundancy-exhaustion story because
the breakpoint is predicted before training.

## Auxiliary Task Choice: Low-Rank Linear Teacher

Use:

\[
x\sim \mathcal N(0,I_{n_0}),
\]

\[
y=A x + \epsilon,
\]

where:

\[
A=U\Sigma V^\top,\qquad \operatorname{rank}(A)=r.
\]

Reasons:

- Input/output dimensions are explicit.
- Task rank is controlled.
- Test loss is easy to measure.
- It is simple enough that BP should establish a clear capacity baseline.

Open risk:

- The task may be too easy for ReLU MLPs, and FA may not show a clean transition unless width is very small.

## Architecture

Use a two-hidden-layer ReLU MLP:

\[
n_0 \to n \to n \to n_L.
\]

Sweep hidden width:

\[
n\in\{2,3,4,6,8,12,16,24,32,48,64\}.
\]

Parameter count without bias:

\[
P(n)=n_0 n+n^2+n n_L.
\]

Depth is fixed so that the main knob is capacity, not feedback path length.

## Training

For each width:

- Train BP baseline.
- Train FA with fixed random feedback matrices.
- Use same initial forward weights for matched BP/FA comparisons.
- Use multiple data/init seeds and multiple feedback seeds.

Initial smoke sweep:

- widths: \(\{2,4,8,16,32\}\)
- init seeds: `3`
- feedback seeds per init: `5`

Scale-up sweep:

- widths: `2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64`
- init seeds: `10`
- feedback seeds per init: `20`

## Theoretical Quantities to Estimate

### Parameter Count

\[
P=\sum_l n_l n_{l-1}.
\]

### Task Effective Dimension

At the BP-trained endpoint, estimate:

\[
J=\frac{\partial f_\theta(X)}{\partial \theta}
\]

on a probe batch. Let:

\[
K=JJ^\top.
\]

Effective dimension:

\[
d_{\mathrm{eff}}(\lambda)
=
\operatorname{tr}\left[K(K+\lambda I)^{-1}\right].
\]

This estimates how many functional directions the task uses locally.

### Hard FA Burden

For the random-label memorization task, use the hard alignment-burden proxy:

\[
K_{\mathrm{FA}}=\sum_l(D_l-1).
\]

For ordinary FA in a two-hidden-layer MLP:

\[
D_1=n^2,\qquad D_2=n n_L,
\]

so:

\[
K_{\mathrm{FA}}=(n^2-1)+(n n_L-1).
\]

### Log-Volume FA Burden

For matrix-level capacity validation, keep the log-volume quantity:

\[
k_{\mathrm{FA}}^{\mathrm{proxy}}
=
\sum_l C_l(q),
\]

where:

\[
C_l(q)=
-\log P(Q_l\ge q),
\qquad
Q_l\sim \mathrm{Beta}\left(\frac12,\frac{D_l-1}{2}\right).
\]

Use a fixed threshold \(q\) for the first sweep. This is not a hard constraint
count; it is the matrix-level log-volume cost used in the direct capacity
distribution experiments.

### Redundancy Score

Define:

\[
R_{\mathrm{eff}}
=
P-d_{\mathrm{eff}}-K_{\mathrm{FA}}.
\]

For random-label memorization, also use the direct theoretical margin:

\[
M_{\mathrm{FA}}=P-K_{\mathrm{FA}}-Nn_L.
\]

## Main Plots

1. BP and FA train MSE vs hidden width or parameter count.
2. FA/BP train MSE gap vs hidden width or parameter count.
3. FA/BP train MSE gap vs predicted FA margin \(P-K_{\mathrm{FA}}-Nn_L\).
4. Estimated \(d_{\mathrm{eff}}\), \(k_{\mathrm{FA}}^{\mathrm{proxy}}\), and \(P\) vs width.

## Success Criterion

The experiment supports the contribution if:

- BP reaches low train MSE near the predicted \(P\gtrsim Nn_L\) threshold.
- FA reaches low train MSE near the predicted \(P-K_{\mathrm{FA}}\gtrsim Nn_L\) threshold.
- The FA/BP train gap is largest where BP has positive margin but FA has negative margin.
- The transition aligns better with \(P-K_{\mathrm{FA}}-Nn_L\) than with arbitrary width.

Teacher-regression tasks are useful only as auxiliary downstream checks because
their task dimension is not known before training.