Repository navigation
Description
Summary
simplexity.analysis.linear_regression calls jnp.linalg.lstsq(..., rcond=None). With
rcond=None, numpy and jax use a cutoff of max(M, N) * eps on the singular values of the design —
scaled by the number of rows. Fitting activations, M is the number of samples, so the cutoff
grows with the dataset: the more data you fit on, the more of the design is thrown away. Under
jax_enable_x64 = False the relevant eps is float32's, making the effect large at ordinary sizes.
Measurement
Regressing a transformer's residual stream onto a 131-dimensional target, 66,930 training rows,
28,869 held-out rows (experiments/toki_pona/targets/ in simplex-research):
| held-out R² | singular directions kept | |
|---|---|---|
linear_regression (as shipped) |
0.609235 | 56 of 321 |
np.linalg.lstsq float64, same design |
0.740682 | 321 of 321 |
| moment-based OLS, independent implementation | 0.740679 | 320 of 321 |
The design's condition number is 5.3e3 — it needs no truncation at all. The cutoff is
66930 * 1.1921e-7 = 7.98e-3 of the largest singular value, so every direction below 0.8% of
the top one is zeroed. 265 well-conditioned directions are discarded, and the held-out fit is
0.13 R² worse for it.
tp-seeds separated the two candidate causes by running float32 with an explicit rcond: it
agrees with float64 to six decimals. So the mechanism is the cutoff, not the precision.
Why it is worse than a constant bias — it moves an argmax
The number of directions above the cutoff depends on the design's spectrum, which varies by site
and by checkpoint. Measured over one capacity ladder, the directions kept at the worst site are
30 of 192 (small), 59 of 320 (ref) and 17 of 512 (large), while at the site after the final
LayerNorm — which flattens the spectrum — essentially all of them survive and the two solvers agree
to 8e-4.
So the bias is a function of depth, and the harm is not that a curve sits low. It is that a
curve read along depth has the wrong shape: on this ladder the belief's decodability peaks at the
first or second block in float64 and at ln_final under the library, so the reported peak layer
is wrong, not merely the value at it. A second session reproduced the same argmax move
independently on a different model family.
And no summary a caller can see predicts how much is discarded. Over 62 sites of nine
checkpoints the fraction of directions surviving the cutoff ranges from 3.3% to 100%.
- Width does not predict it. Two checkpoints of the same architecture (d_model 192) keep 193 of
193 at every site and 30 of 192 respectively — a fifty-fold difference in the resulting error. - The condition number predicts it only loosely. Spearman between
log10(condition)and the
fraction kept is −0.68 over those 62 sites, and individual pairs invert: a second session
measured a well-conditioned site (κ = 485) losing twelve directions while a near-singular one
(κ = 6.6e7) on the same model lost two, because a fixed relative cutoff sweeps away whatever
happens to sit in the 2e-3…8e-3 band.
The amount discarded is an uncontrolled functional of the whole spectrum. A caller cannot bound
their exposure without running the SVD themselves — which is the thing they called this function to
avoid.
And the error is not one-signed — it is an unchosen regulariser
Truncating singular values is regularisation, so where the design is genuinely ill-conditioned it
helps. At large's embedding site, condition number 2.25e9:
| held-out R² | directions kept | |
|---|---|---|
linear_regression (as shipped) |
0.5885 | 232 of 512 |
| plain least squares, float64 | −0.0413 | 512 of 512 |
Plain OLS is worse than predicting the mean there, and the library's cutoff is the only thing
preventing the blow-up. So the defect is not "the numbers come out low". It is that
linear_regression applies a regularisation strength nobody chose, which varies by dataset size,
by site and by checkpoint, and is sometimes far too strong and sometimes exactly what saves the
fit. A caller who wants least squares does not get it; a caller who wants regularisation cannot
set it; and neither can tell which they got.
The remedy a caller needs is therefore not only an explicit rcond but a stated regulariser. On
the measurement above, a ridge chosen on a validation split dominates both of these estimators at
every site, and costs nothing when the problem is solved in moments.
Suggested fix
Pass an explicit rcond (or solve in float64). rcond=None scaling by row count is a long-standing
numpy wart; the documented remedy is to set it yourself. Two further things would help a caller:
- Return the effective rank alongside the coefficients. A caller cannot currently tell that
82% of their design was dropped. - Offer a held-out scorer.
linear_regressionreports metrics on the data it was fitted to;
every caller that wants an out-of-sample R² reimplements it, and they will not all agree on the
weighting. (stage2/README.mdalready names this gap.)
Provenance. Measured independently by two sessions on two checkpoint families; write-up and reproduction in simplex-research at experiments/toki_pona/targets/simplexity-issue.md (branch eric/toki-pona, a5d401baa), with the ridge comparison in experiments/toki_pona/targets/README.md §2. Filed by jan-sona on behalf of Eric Alt.
Activity