Scale Determines Whether Language Models Organize Representation Geometry for Prediction
For researchers studying language model representations, this reveals a scale-dependent relationship between geometry and prediction that loss curves and spectral metrics fail to capture.
The paper introduces Subspace PGA to measure whether language model representation geometry aligns with the unembedding readout subspace. It finds that intermediate layers are significantly organized for prediction (peak z=9-24), but small models (d≤1024) progressively lose this alignment at late layers during training, while large models (d≥2048) preserve it.
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 $W_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 = 9$--$24$), but the degree is scale-dependent: small models ($d \leq 1024$) progressively lose it at late layers during training -- even as loss keeps improving -- while large models ($d \geq 2048$) preserve it throughout. We trace this to a capacity trade-off: a few dominant directions migrate away from $W_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.