💬 NLP

Disentangling Geometry, Performance, and Training in Language Models

This paper systematically investigates the relationship between Transformer unembedding matrix geometry and language model performance across 108 controlled models, concluding that geometric metrics like effective rank primarily reflect training hyperparameters rather than serving as reliable predictors of downstream performance.

Original authors: Atharva Kulkarni, Jacob Mitchell Springer, Arjun Subramonian, Swabha Swayamdipta

Published 2026-02-25

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Atharva Kulkarni, Jacob Mitchell Springer, Arjun Subramonian, Swabha Swayamdipta

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

Imagine you are a chef trying to figure out why some dishes turn out delicious and others taste like cardboard. You decide to look at the kitchen tools (the geometry of the weights) to solve the mystery. Specifically, you're looking at a special tool called the "unembedding matrix," which is like the final plating station where the chef decides what word comes next.

For a while, other chefs (researchers) believed that if this plating station looked "flat" or "collapsed" (low effective rank), the dish would be bad. They thought a "flat" tool meant the food was ruined.

This paper is like a team of food scientists who decided to test that theory by cooking 108 different versions of the same recipe, changing everything from the size of the pan (batch size) to how much salt they used (weight decay). Here is what they found, explained simply:

1. The "Flat Tool" Myth

The Old Belief: "If the plating station looks squashed (low rank), the food is bad."
The New Finding: Not necessarily!
Think of the "effective rank" as the diversity of flavors on the plate. The scientists found that while a diverse plate often tastes good, a flat plate doesn't always taste bad.

The Twist: Sometimes, a model with a "flat" tool (low rank) still cooks a delicious meal. Other times, a model with a "diverse" tool (high rank) serves a terrible dish.
The Analogy: It's like having a fancy, multi-bladed knife (high rank). Just because you have a fancy knife doesn't mean you'll make a great steak. Conversely, a simple butter knife (low rank) can still cut a perfect steak if you know how to use it. The tool's shape isn't the magic ingredient; how you use it is.

2. The Real Culprit: The Recipe, Not the Tools

The scientists discovered that the shape of the tool (the geometry) is actually just a reflection of the recipe choices (hyperparameters), not the quality of the food itself.

Batch Size (How much you cook at once): If you cook in huge batches, your tool tends to stay "diverse" (high rank). If you cook in tiny batches, it tends to get "squashed" (low rank). But cooking in huge batches doesn't guarantee a better taste; sometimes it makes the food worse if you don't adjust the heat (learning rate).
Weight Decay (How much you restrict the ingredients): Adding more "restriction" (regularization) keeps the tool diverse. But again, this doesn't automatically make the food taste better. It just changes the shape of the tool.

The Big Takeaway: The shape of the tool is a symptom of how you cooked, not the cause of how it tastes. If you see a "flat" tool, don't panic and throw away the dish. It might just mean you used a specific cooking method.

3. The "Burnt Toast" Problem (Saturation)

There was a famous theory that small models eventually "burn out" (saturation) because their tools get too flat.

The Experiment: The team cooked a small model using a different oven (architecture) than the one that usually burns.
The Result: Even though the tool got very flat, the food didn't burn. The model kept tasting good even after being cooked for a very long time.
The Lesson: The "burning" wasn't caused by the flat tool. It was caused by the specific type of oven (the architecture) or the heat settings. So, low rank is a side effect, not the arsonist.

4. Other Tools Don't Help Much

The scientists also looked at other ways to measure the kitchen tools, like checking how similar the flavors are to each other (Cosine Similarity) or how evenly they are spread out (Isotropy).

The Finding: These other measurements were just echoing the same story as the "flatness" measurement. They didn't give any new clues about whether the food would taste good. They were like looking at the color of the knife to predict the taste of the soup—it just doesn't work.

Summary: What Should You Do?

If you are building or training AI models, here is the practical advice from the paper:

Don't judge a book by its cover (or a model by its geometry): If you see a model with a "low rank" (flat) unembedding matrix, don't assume it's broken. It might still be a champion.
Focus on the recipe, not the tools: Instead of trying to force the model to have a specific shape, focus on tuning your batch size, learning rate, and weight decay. These are the knobs that actually control the taste.
Geometry is a diagnostic, not a crystal ball: Use these geometric measurements to understand how the model learned (e.g., "Oh, we used a huge batch size"), but don't use them to predict how well the model will perform on a new task.

In a nutshell: The shape of the AI's brain is a mirror reflecting how it was trained, not a crystal ball predicting how smart it will be. To get a good AI, focus on the training process, not the shape of the weights.

1. Problem Statement

Recent research in language model (LM) interpretability has focused on the geometric properties of Transformer weights, particularly the unembedding matrix ( $W$ ), which maps hidden states to vocabulary logits. A prevailing hypothesis, largely derived from studies on the Pythia model family, suggests that a low effective rank in the unembedding matrix (indicating a collapse of the singular value spectrum) is a primary cause of performance degradation, specifically in small models undergoing "saturation" (late-stage training performance drops).

However, the authors identify critical gaps in this narrative:

Scope Limitation: Previous findings are based almost exclusively on the Pythia model family.
Causality vs. Correlation: Prior work does not systematically vary training conditions to determine if low effective rank causes poor performance or merely co-occurs with other confounding factors (e.g., skewed token distributions, optimization instability, or specific hyperparameter choices).
Metric Ambiguity: It is unclear if other geometric metrics (isotropy, cosine similarity) or final-layer representations provide better predictive power than effective rank.

The core research questions are:

RQ1: Does the effective rank of the unembedding matrix reliably predict model performance across diverse tasks and training configurations?
RQ2: Why does effective rank often fail to correlate with performance?
RQ3: How do alternative geometric metrics and learned representations relate to performance and effective rank?

2. Methodology

To address these questions, the authors conducted a large-scale, controlled empirical study involving 108 OLMo-style language models.

Model Suite: Models ranging from 4M to 75M non-embedding parameters, trained on the Pile dataset.
Training Variables: The study systematically ablated key hyperparameters to isolate their effects:
- Token Budgets: 8B, 32B, 64B, and 128B tokens.
- Batch Sizes: 32, 64, 128, 256, 512.
- Weight Decay: 0.01, 0.05, 0.1.
- Learning Rate (LR): Scaled relative to model-specific optimal rates (0.1x, 10x, 100x).
- LR Annealing: Decay to 0%, 1%, and 10% of peak LR.
Geometric Metrics Analyzed:
- Effective Rank ( $R(W)$ ): Based on the entropy of the singular value spectrum (Roy & Vetterli, 2007).
- Isotropy ( $I(W)$ ): Measures uniformity of vector distribution (Arora et al., 2016).
- Angular Variability / Cosine Similarity ( $A(W)$ ): Measures alignment between vectors (Ethayarajh, 2019).
- Representations: Metrics were also computed for the final-token's last-layer representations ( $H$ ).
Evaluation Tasks:
- In-Distribution (ID) and Out-of-Distribution (OOD) pretraining loss.
- Fine-tuning performance (StarCoder-Python, OpenWebMath).
- Catastrophic forgetting (performance retention on pretraining data after fine-tuning).
- Post-training quantization (PTQ) robustness.
- Saturation Analysis: Direct comparison between OLMo-14M and Pythia-14M to investigate late-stage degradation.

3. Key Contributions & Results

A. Effective Rank is Correlated but Not Deterministic (RQ1)

General Trend: Higher effective rank generally correlates with lower loss (better performance), but this is not a guarantee.
Counterexamples: The authors found numerous models with high effective rank that performed poorly, and models with low effective rank that maintained good performance.
Task Dependency:
- Pretraining: ID loss tracks effective rank loosely, but OOD loss tracks ID loss much more reliably than it tracks effective rank.
- Quantization: Low effective rank is strongly associated with severe performance degradation after quantization, particularly in smaller models.
- Fine-tuning & Forgetting: Neither effective rank nor pretraining loss reliably predicts fine-tuning performance or catastrophic forgetting. Fine-tuning performance appears largely initialization-invariant regarding geometry.

B. Hyperparameters Drive Geometry, Not Performance (RQ2)

The study reveals that geometric metrics are primarily symptoms of training dynamics rather than causes of performance.

Batch Size: Larger batch sizes consistently preserve higher effective rank. However, performance follows a U-shaped curve relative to batch size; extremely large batches can degrade performance despite high rank.
Weight Decay: Stronger weight decay preserves high effective rank. However, its impact on performance is model-size dependent (small models prefer weak regularization; larger models benefit from moderate regularization).
Learning Rate Annealing: Aggressive annealing (decaying LR to 0) improves performance but has little effect on rank. Mild annealing (keeping LR at 10%) preserves high rank but leads to higher loss.
Conclusion: The observed correlation between rank and performance is largely confounded by hyperparameter choices.

C. Saturation is Not Caused by Low Rank (RQ1 & RQ2)

Pythia vs. OLMo: The authors replicated the Pythia-14M configuration using OLMo architecture.
- Pythia-14M: Exhibited late-stage loss degradation coinciding with a sharp drop in effective rank.
- OLMo-14M: Trained under identical conditions, showed no loss degradation even when effective rank collapsed significantly (especially with small batch sizes).
Implication: Low effective rank is neither necessary nor sufficient to explain saturation. Saturation in Pythia may be an architectural artifact (e.g., parallel attention mechanisms) or specific optimization instability, not a universal geometric phenomenon.

D. Alternative Metrics Offer Limited Insight (RQ3)

Cosine Similarity: Shows a non-monotonic ("moon-shaped") relationship with effective rank. It does not provide additional predictive power beyond effective rank.
Isotropy: Remains largely uncorrelated with performance and varies minimally across different training setups, except for a strong dependence on weight decay.
Representations ( $H$ ): The geometry of final-layer representations is more stable and less variable than the unembedding matrix ( $W$ ). They do not correlate well with downstream performance differences.

4. Significance and Practical Takeaways

This paper fundamentally shifts the perspective on model geometry from a predictive proxy to a diagnostic tool for training behavior.

Do Not Use Geometry for Early Stopping: Halting training because effective rank is decreasing is ill-advised. Low rank can coexist with high performance, and high rank does not guarantee success.
Hyperparameters > Geometry: Practitioners should focus on tuning standard hyperparameters (batch size, weight decay, token budget) rather than attempting to optimize geometric properties directly. Geometry is a reflection of these choices.
Saturation is Contextual: The "saturation" phenomenon observed in small Pythia models is not a universal law of small LMs. Architectural choices (e.g., sequential vs. parallel attention) play a critical role.
Quantization Sensitivity: While geometry doesn't predict general performance well, low effective rank is a strong indicator of fragility to post-training quantization.
Causality Warning: The authors emphasize that current geometric metrics are observational. Future work requires interventional studies (manipulating rank directly) to establish causality.

Conclusion

The authors conclude that the geometry of the unembedding matrix, as captured by effective rank and related metrics, primarily reflects optimization dynamics shaped by training hyperparameters rather than intrinsic model performance capabilities. Consequently, these metrics should be used to understand how a model was trained, not to predict how well it will perform on downstream tasks.