Disentangling Geometry, Performance, and Training in Language Models

TL;DR

The findings suggest that the model's geometry, as captured by existing metrics, primarily reflects training choices rather than performance, which is in contrast to prior work.

Abstract

Geometric properties of Transformer weights, particularly the unembedding matrix, have been widely useful in language model interpretability research. Yet, their utility for estimating downstream performance remains unclear. In this work, we systematically investigate the relationship between model performance and the unembedding matrix geometry, particularly its effective rank. Our experiments, involving a suite of 108 OLMo-style language models trained under controlled variation, reveal several key findings. While the best-performing models often exhibit a high effective rank, this trend is not universal across tasks and training setups. Contrary to prior work, we find that low effective rank does not cause late-stage performance degradation in small models, but instead co-occurs with it; we find adversarial cases where low-rank models do not exhibit saturation. Moreover, we show that effective rank is strongly influenced by pre-training hyperparameters, such as batch size and weight decay, which in-turn affect the model's performance. Lastly, extending our analysis to other geometric metrics and final-layer representation, we find that these metrics are largely aligned, but none can reliably predict downstream performance. Overall, our findings suggest that the model's geometry, as captured by existing metrics, primarily reflects training choices rather than performance.

Figures (27)

• Figure 1: (A) The effective rank of a model's unembedding matrix, $\boldsymbol{\mathrm{W}}$, tends to correlate with its performance, but does not guarantee it. (B) The effective rank of the unembedding matrix is influenced by hyperparameter choices (e.g. batch size). (C) These hyperparameters in turn can affect model performance. (D) The effective rank of the last token's final-layer representation, $\boldsymbol{\mathrm{H}}$ shows little correlation with that of the unembedding matrix and varies minimally across different hyperparameters (such as batch size) (B). Note: The green dots are used to produce the line plots. The loss is normalized within each model size by its lowest observed value.
• Figure 2: In-distribution evaluation: Variation in effective rank does not translate to equivalent change in in-domain loss.
• Figure 3: Out-of-Distribution Evaluation: Performance in OOD setup remains largely unchanged across different effective ranks $\mathcal{R}(\boldsymbol{\mathrm{W}})$. Alternatively, models with stronger in-domain loss consistently transfer well, making it a better predictor of OOD behavior.
• Figure 4: Post-training quantization evaluation: Highly collapsed models are heavily impacted, while high-rank models remain largely intact but not entirely consistent.
• Figure 5: Saturation evaluation: Not all small language models experience saturation. Low effective rank alone does not cause it. Instead, saturation may result from other training or architectural factors that can also induce rank collapse.