💻 computer science

Scale Determines Whether Language Models Organize Representation Geometry for Prediction

This paper introduces Subspace PGA to reveal that while language model representation geometry is organized for prediction, small models lose this alignment in later layers due to a capacity trade-off masked by dominant directions, whereas large models preserve it throughout training.

Original authors: Weilun Xu

Published 2026-05-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Weilun Xu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Idea: How Language Models "Think" in Space

Imagine a language model (like the ones that write essays or chat with you) as a giant, multi-story library. Every time the model reads a word, it creates a "bookmark" (a representation) and places it on a specific shelf in a specific room (a layer) of the library.

The paper asks a simple question: How are these bookmarks arranged?

For a long time, scientists looked at the shape of the shelves (are they crowded? are they empty?). This paper asks something different: Are the shelves arranged to help the model predict the next word?

The authors discovered that the answer depends entirely on the size of the library (the model's scale).

The Tool: The "Prediction Compass" (Subspace PGA)

To answer this, the authors invented a new measuring tool called Subspace PGA.

Think of the model's final goal as a "Target Zone" (the unembedding matrix, or $W_U$ ). This is the specific direction the model needs to point to in order to guess the next word correctly.

The Test: The authors take the "bookmarks" from the middle of the library and ask: "If we only look at the directions that point toward the Target Zone, do the bookmarks still make sense relative to each other?"
The Score (Z-score):
- High Positive Score: The bookmarks are perfectly organized. Even if you ignore all the other directions, the ones pointing to the Target Zone still hold the map together. The library is organized for prediction.
- Zero or Negative Score: The bookmarks are scattered. The directions pointing to the Target Zone are just as random as any other direction. The library has lost its organization for prediction.

The Discovery: Size Matters

The researchers tested libraries of different sizes (from tiny 70-million-parameter models to huge 6.9-billion-parameter models). They found two very different behaviors:

1. The Large Libraries (Large Models)

The "Direct Route"
In big models (like Pythia-1B and larger), the organization is consistent from the bottom floor to the top.

Analogy: Imagine a massive highway system. No matter which exit you take, the road signs are always clear and point directly toward the destination. The "Target Zone" directions are always the most important ones.
Result: The geometry stays organized for prediction all the way through.

2. The Small Libraries (Small Models)

The "Detour"
In smaller models (like Pythia-410M and smaller), something strange happens near the top floors (the later layers).

Analogy: Imagine a small, crowded town. As you get closer to the exit, the main road gets blocked by a massive, slow-moving truck (a "dominant direction" of variance). This truck isn't going toward the Target Zone; it's going somewhere else.
The Problem: Because the town is small, there isn't enough room for both the "Target Zone" road and the "Truck" road. The Truck takes over the whole street. The signs pointing to the Target Zone get pushed aside and become invisible.
The Twist: Just before the very last exit (the final layer), the model suddenly clears the truck away and puts the signs back up. The final answer is still correct, but the journey there was a detour.

Key Findings in Plain English

1. The "Masking" Effect
The paper argues that the small models aren't forgetting how to predict. The information is still there, but it's masked (hidden) by that one giant "Truck" direction.

Proof: If you artificially remove that one dominant direction (the truck), the small models suddenly show perfect organization again, just like the big ones. The structure was never destroyed; it was just covered up.

2. Loss Curves Lie
Usually, we judge a model by its "Loss" (how many mistakes it makes).

The Surprise: Both the small models (with the detour) and the big models (with the direct route) get the same low error rate. They both predict the next word equally well.
The Lesson: You cannot tell if a model is "organized" or "detoured" just by looking at its error rate. A model can be making perfect predictions while its internal geometry is completely scrambled.

3. Old Tools Missed This
Scientists previously used tools that measured the "shape" of the data (like how spread out the bookmarks are).

The Failure: These old tools said the small models and big models looked the same. They couldn't see that the small models were taking a detour. The new "Prediction Compass" (Subspace PGA) is the first tool to spot this difference.

Summary Analogy

Imagine two chefs making the same soup.

The Big Chef (Large Model): Keeps all the ingredients organized in labeled bowls from start to finish. When it's time to serve, the soup is ready.
The Small Chef (Small Model): Starts with labeled bowls. But halfway through, they dump everything into one giant, chaotic pot because their kitchen is too small to keep the bowls separate. They stir it around wildly (the "detour"). However, right before serving, they magically scoop out the perfect soup from the chaos and serve it.

The Result: Both chefs serve the same delicious soup (low loss).
The Difference: The Big Chef's kitchen was organized for the task the whole time. The Small Chef's kitchen was chaotic and disorganized for most of the process, only fixing itself at the very last second.

The Conclusion: The size of the model determines not just how well it predicts, but how it organizes its internal world to do so. Small models take a chaotic detour; large models take a direct path.

Technical Summary: Scale Determines Whether Language Models Organize Representation Geometry for Prediction

Problem Statement

In language models (LMs), the semantic content of a representation is determined by the geometry of its representation space—specifically, the distances between points rather than their absolute activations. While existing tools characterize the shape of this geometry (e.g., rank, spectral decay, cross-model convergence), they do not address the functional question: what is this geometry organized for?

Theoretical frameworks suggest that representations should organize around the model's predictive task, specifically the unembedding matrix $W_U$ which maps hidden states to the next-token distribution. However, it remains an empirical question whether intermediate layers align their distance structures with this readout subspace. Current metrics, such as the Logit Lens (which tests decodability of individual points) or spectral metrics (which describe manifold shape), fail to capture whether the geometric arrangement of distances is functionally oriented toward prediction. Furthermore, it is unknown if this organization is maintained uniformly across all layers and model scales.

Methodology: Subspace PGA

The authors introduce Subspace Predictive-Geometric Alignment (Subspace PGA), a metric designed to quantify the degree to which a layer's distance structure aligns with the model's predictive readout.

Readout Subspace Definition: The unembedding matrix $W_U$ is decomposed via Singular Value Decomposition (SVD). The top- $k$ right singular vectors define the readout subspace ( $V_k$ ), representing the directions in hidden space that $W_U$ amplifies most to produce predictions.
Distance Preservation: For a given layer $\ell$ , the authors compute pairwise cosine distances between hidden states in the full $d$ -dimensional space ( $d_{full}$ ) and in the projected readout subspace ( $d_{proj}$ ).
Null Distribution: To isolate the specific contribution of the readout, the authors compare the Spearman correlation between $d_{full}$ and $d_{proj}$ against a null distribution. This null is generated by projecting the same hidden states onto 100 random $k$ -dimensional subspaces sampled uniformly from the Grassmannian.
Z-Score Calculation: Subspace PGA is defined as the z-score:
$PGA(\ell) = \frac{\rho_{readout} - \mu_{null}}{\sigma_{null}}$
A high positive $z$ indicates the geometry is organized for prediction (the readout preserves more structure than random subspaces). A $z \approx 0$ implies no special alignment, and $z < 0$ implies the readout is less aligned with the distance structure than random subspaces.
Anisotropy Correction: To prevent the dominant variance direction (which inflates cosine similarity) from confounding the metric, the authors apply mean-centering followed by CCR-1 (removing the leading principal component) before computing distances.

Key Results

The study evaluates seven Pythia models (70M to 6.9B parameters) and three cross-family models (OLMo, Phi, Gemma).

1. Scale-Dependent Organization

Large Models ( $d \ge 2048$ ): These models maintain strong predictive organization throughout all intermediate layers. The $z$ -scores remain significantly positive (peaking at $z=9–24$ in mid-layers), indicating that the distance structure is consistently aligned with the readout subspace.
Small Models ( $d \le 1024$ ): These models exhibit a "detour" regime. While they establish predictive organization early in training and in early layers, they progressively lose this alignment in late-but-not-final layers (e.g., layers 20–23 in Pythia-410M). In these "collapse zones," $z$ -scores drop significantly (e.g., to $-32$), indicating that the distance structure becomes misaligned with the readout, even as the model's validation loss continues to improve.
Recovery: In small models, the final layer recovers the alignment, projecting the geometry back into the readout subspace before the output.

2. Mechanism: Capacity Trade-off and Masking

The loss of organization in small models is not a destruction of predictive information but a masking effect caused by capacity constraints.

Dominant Direction Migration: In the late layers of small models, the direction of largest variance in the residual stream (PC1) rotates away from the readout subspace into "dark" directions (orthogonal to $W_U$ ).
Masking: Because the dominant variance direction carries most of the distance structure, random subspaces (which capture a fraction of this variance) end up preserving the distance structure as well as, or better than, the readout subspace. This causes the $z$ -score to drop.
Restoration: Removing the top few principal components (CCR-5 or CCR-10) restores positive $z$ -scores in all layers of models with $d \ge 768$ . This confirms the predictive structure was intact but obscured by a few dominant, prediction-irrelevant directions.

3. Independence from Spectral Metrics and Loss

Spectral Blindness: Standard spectral metrics (RankMe, $\alpha$ -ReQ, isotropy, condition number) fail to detect this distinction. In the collapse zones of small models, spectral metrics remain stable or flat while Subspace PGA swings by over 40 standard deviations.
Loss Insensitivity: Validation loss curves are nearly identical across small and large models, failing to capture the geometric reorganization occurring in intermediate layers.

4. Cross-Architecture Convergence

The readout subspace concentrates cross-architecture convergent structure. When comparing different models (e.g., Pythia vs. OLMo), the Spearman correlation of pairwise distances is higher in the 100-dimensional readout subspace than in the full space or random subspaces. This suggests that what independently trained models agree on geometrically is specifically tied to the predictive function.

5. Downstream Utility

In large models, layers with high Subspace PGA $z$ -scores also yield higher accuracy for linear probes on downstream tasks (AG News). In small models, probe accuracy remains high even in the "collapse zone" where $z$ is negative, further proving that the underlying features are preserved but merely misaligned with the readout interface.

Significance and Claims

The paper claims that scale determines not only how well a model predicts, but how its representation geometry is organized to do so.

Two Geometric Regimes: The authors identify a "direct" regime (large models) where geometry stays aligned with prediction, and a "detour" regime (small models) where late layers temporarily route computation through prediction-irrelevant directions before correcting at the output.
Functional Orientation: The study shifts the focus from characterizing the shape of representation geometry to understanding its functional orientation. It demonstrates that geometry can be highly structured yet functionally misaligned with the model's objective.
Limitations of Current Metrics: The work highlights that loss curves and spectral metrics are insufficient for diagnosing the functional organization of intermediate representations. Subspace PGA provides a task-free, dimensionality-matched metric to detect these phenomena.
Capacity vs. Architecture: The phenomenon is attributed to representational capacity (width) rather than architectural specifics, as it is observed across different model families when dimensions are small.

The authors conclude that the "detour" in small models is a side effect of limited capacity forcing a collision between intermediate computation directions and readout-aligned directions, rather than an explicitly learned routing strategy. This distinction is invisible to standard training signals but critical for understanding the internal mechanics of language models.