English

Scale Determines Whether Language Models Organize Representation Geometry for Prediction

Machine Learning 2026-05-19 v1 Computation and Language

Abstract

In language models, what a representation encodes is determined by the geometry of its representation space: distances, not activations, carry meaning. Existing tools characterize the shape of this geometry but do not ask what that shape is organized for. We introduce Subspace PGA, a metric that tests whether a layer's distance structure aligns with the readout subspace of the unembedding matrix WUW_U more than with random subspaces of equal size. Across seven Pythia models (70M--6.9B) and three cross-family models, intermediate geometry is significantly organized for prediction (peak z=9z = 9--2424), but the degree is scale-dependent: small models (d1024d \leq 1024) progressively lose it at late layers during training -- even as loss keeps improving -- while large models (d2048d \geq 2048) preserve it throughout. We trace this to a capacity trade-off: a few dominant directions migrate away from WUW_U's readout, masking rather than destroying the predictive structure beneath, and removing them restores alignment. Neither spectral metrics nor loss curves capture this distinction. Scale thus determines not only how well a model predicts, but how its representation geometry is organized to do so.

Keywords

Cite

@article{arxiv.2605.17084,
  title  = {Scale Determines Whether Language Models Organize Representation Geometry for Prediction},
  author = {Weilun Xu},
  journal= {arXiv preprint arXiv:2605.17084},
  year   = {2026}
}