Skip to content

analysis.linear_regression: lstsq rcond=None silently truncates the design with a row-count-scaled cutoff, moving peak layers #201

Open

Description

@ealt

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_regression reports 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.md already 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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions