Tracing the Representation Geometry of Language Models from Pretraining to Post-training

We thank the reviewer for their encouraging comments and suggestions, and we are glad that they found our results to be novel, rigorous and clearly presented. Below, we address all the concerns and questions raised by the reviewer.

Results on complex reasoning tasks

We thank the reviewer for this suggestion. We have now conducted downstream task evaluations for checkpoints from intermediate pretraining stages on Lambada_openai, MMLU, and HellaSwag, and will add them in the appendix of our final paper. It is important to note that models solely undergoing pretraining often struggle with complex reasoning tasks, requiring specific post-training methods to achieve high performance. Therefore, their performance on such tasks is not a reliable indicator of their underlying capabilities. Consequently, Section 3.4, which presents results from post-training stages, focuses on evaluating model performance on complex tasks like chat (AlpacaEval) for SFT models and mathematical reasoning (AMC-23) for RLVR models.

Utility of RankMe and $\alpha$-ReQ for intermediate layer representations

While we focus on and present the results for the final layer representations, the spectral metrics that are used to quantify the representation geometry are generic and do not assume anything about the origin of the representations. Therefore, they can be readily applied to study changes in representation geometry in intermediate layers. To further illustrate this point, we performed experiments on the intermediate layer representations and observed similar trends. We will add these results in order to clarify the utility of the proposed metrics in the Appendix in the final version of our paper.

Perturbing eigenvectors for better out-of-distribution generalization

The relationship between better out-of-distribution generalization and flat minima implies a particular conditioning of the Hessian of the loss with respect to the model parameters. While the connection between this Hessian matrix and the eigenspectrum of the final layer representations is not immediately obvious, it becomes clear in cases where only the final readout weight matrix is adjusted for a new task. In such cases, perturbing specific eigenvectors can improve the optimization problem and, consequently, out-of-distribution generalization. However, more theoretical and empirical work is needed to fully understand the relationship between the geometry of final layer representations and a model's out-of-distribution performance after pretraining. We hope our proposed metrics and observed trends can inspire future strategies and regularization techniques to enhance out-of-distribution generalization.

Perturbing eigenvectors to alter downstream performance

This is an interesting question, highlighting the implications of our findings for mechanistic interpretability. While a full investigation of the causal relationship between the eigenspectrum of representations and model performance is beyond the scope of this work, we believe that our analytical results point to potential causal interventions. Specifically, we show that high-frequency token information is learned early in training and is encoded in the top eigendirections. Consequently, removing these directions would have minimal impact on downstream question-answering task performance. Conversely, preserving the top eigendirections while removing the tail of the eigenspectrum would retain most high-frequency token information but discard fine-grained dataset details, thereby negatively affecting downstream performance.

Given the time constraints of the NeurIPS review process, we will defer a more rigorous analysis of how downstream task performance connects to this information segregation between the top eigendirections and the eigenspectrum tail to future work. Instead, we will add these preliminary results in the Appendix. These results will demonstrate the information segregation and the differential roles of the top eigendirections and the eigenspectrum tail on specific question-answering tasks, illustrating potential causal links between the eigenspectrum and the model's downstream behavior.

We thank the reviewer once again for their overall positive review and are grateful for their helpful suggestions. We believe that adding the results on more diverse tasks during different phases in the pretraining stage, trends in representation geometry changes in intermediate representations and preliminary experiments demonstrating the effect of ablating the eigenspectrum tail on downstream task performance in the appendix will strengthen our paper and increase its impact. We hope that the reviewer finds our responses to be helpful, and are happy to clarify any further questions during the discussion phase.

We thank the reviewer for their helpful comments and suggestions, and we are glad that they found our paper to be well written and the results to be insightful. Below, we address all the concerns and questions raised by the reviewer.

Intermediate layer representations

We agree with the reviewer that our work primarily focuses on the last token of the last layer, and thus, the identified phases are defined specifically based on the representation geometry changes in this representation space. The intermediate layer representations, as pointed out by the reviewer, do not necessarily exhibit the same trends. However, our analysis reveals similar trends in several intermediate layer representations, with the temporal evolution of last layer representations often serving as a leading indicator. This means intermediate layers undergo similar phases, but later in the pretraining stage. We will add these results into the Appendix and clarify the scope of our results in the Discussions section of the final paper.

We defined our focus to be such because, as mentioned in Section 2.1, the last token in the last layer is critical for autoregressive models as: (i) it directly parameterizes the predictive distribution for subsequent tokens, (ii) it synthesizes information from the entire input context, reflecting the model’s contextual understanding, and (iii) it is frequently used as input for task-specific layers in downstream applications. This pivotal role allows us to directly relate spectral metrics to model behavior. But, that is not to say that the intermediate layers are unimportant, and for this reason, we will include those additional results in the updated Appendix.

Spectral measure computation details

We apologize for the lack of clarity in the methodology. We used 15k sequences from the Fineweb dataset and extracted the last layer representations after processing sequences of 512 tokens. We will add these details in Section 2.1 in the final version of our paper. The impact of sequence length and dataset on the spectral metrics is presented in Fig 7 and Fig 8 respectively. We will add references to these figures in Section 2.1 as well. We thank the reviewer for pointing this out.

Section 3.3 experiment details

We appreciate the reviewer's suggestion to add more detail about the analytically tractable experiment, as it will improve the paper's readability. In the experimental setup for Fig 4, our goal is to validate our theoretical findings regarding the role of cross-entropy loss, so we abstract away the complexities of natural language modeling. We define $f_{\theta}(.)$ to be an embedding model that receives inputs $x^{(i)}$'s, which are orthogonal vectors. Each $x^{(i)}$ is randomly assigned to one of four classes (represented by triangle, circle, square, and diamond in Fig 4A). The frequency distribution is skewed: two $x^{(i)}$’s belong to the triangle class, two to the circle class, and one each to the square and diamond classes (as shown in Fig 4A). For the control experiment in Fig 10(a), we use a uniform class distribution, with two samples per class.

Fig 3 plots

We apologize for the lack of clarity in the plots. We will reduce the number of panels in Fig 3 and make the necessary changes as suggested by the reviewer.

We thank the reviewer once again for their comments and are grateful for their helpful suggestions. We believe that adding the results on representation geometry changes in intermediate representations in the appendix and clarifying the scope of our results in the discussions will strengthen our paper. We hope that the reviewer finds our responses to be helpful, and are happy to clarify any further questions during the discussion phase.

We thank the reviewer for their encouraging comments and helpful suggestions, and we are glad that they found our results to be novel, well-supported and theoretically grounded. Below, we address all the concerns and questions raised by the reviewer.

Mechanistic Validation

This is a great question, and forms an important next step to our current work. While we have deferred investigating causal relationships between eigenspectrum of representations and model behavior to future work, we believe that our analytical results shed light into potential causal interventions that can be designed. Specifically, we demonstrate that the high-frequency token information is learned earlier in training and mainly represented in the top eigendirections. Consequently, removing these directions would have minimal impact on downstream question-answering task performance. On the other hand, keeping the top eigendirections and removing the tail of eigenspectrum would imply preserving most of the information about high-frequency tokens and removing fine-grained information about the dataset, thereby hurting downstream performance. While the concepts of memorization and generalization are loosely connected to this information segregation between the top eigendirections and the tail of eigenspectrum, given the time constraints of the NeurIPS review process, we defer a more rigorous analysis to future work. Instead, we propose to add preliminary results demonstrating the information segregation and the differential role of the top eigendirections and the eigenspectrum tail on the task performance in the Appendix, to demonstrate the potential causal links between eigenspectrum and behavior. We thank the reviewer for this suggestion.

Robustness analysis

While we agree with the reviewer's suggestion that it would help to understand the impact of hyperparameter choices, our compute budget limitations prevented us from pretraining large-scale models ourselves. Instead, we used open-source models and checkpoints, demonstrating consistent non-monotonic spectral evolution across Pythia and OLMo suites with varying architectures and scales. Our results are limited to existing hyperparameter configurations of these model families. However, we believe these results encourage future work on understanding the impact of design choices on the multi-phase learning dynamic.

For smaller-scale models, we have preliminary results where we trained GPT-2 models from scratch on Fineweb and observed the emergence of multi-phase learning dynamics across various "good" hyperparameter configurations (monotonic loss decrease, reasonable downstream performance). We will include an Appendix figure illustrating these results. Taken together, we believe that we can broadly say that these phases are present in well-performing LLMs trained with next-token prediction and cross-entropy loss, deferring further investigation into the impact of hyperparameters to future work.

Alternate metrics

While the reviewer is correct in pointing out the existence of various metrics for quantifying representation geometry, we focused on two that rely on linear decomposition. Our choice was motivated by the direct relationship between linear decomposition of the representation space and learning dynamics under gradient descent. This connection allows us to link our findings to empirically observed model behaviors and provide a well-grounded theoretical explanation.

In contrast, metrics like intrinsic dimensionality identify a non-linear transformation of the representation space that explains most of its variance. While useful for understanding the underlying representation manifold, these metrics are not directly connected to gradient descent properties, making it harder to draw concrete conclusions about model behavior.

Furthermore, recent work by Skean et al. [1] (the paper recommended by the reviewer) demonstrates theoretical connections and strong empirical correspondence between several information-theoretic, geometric and spectral metrics in LLMs. Therefore, we believe our results, combined with those of Skean et al., offer a more comprehensive view of representation geometry evolution in LLMs across pretraining and post-training stages. We will add this point in the discussion section of our paper, along with the citation to the recent work from Skean et al.

Layer-wise analysis

We once again thank the reviewer for highlighting the recent work from Skean et al. on layer-wise representation geometry. This interesting work offers a complementary perspective on model behavior, and we will cite it and discuss its connection to our own in the updated discussion section of our paper. Notably, our analysis shows similar trends in intermediate layer representations, with the last layer's temporal evolution often serving as a leading indicator. Specifically, most intermediate layers exhibit similar phases, but later in pretraining. We will add these results in the Appendix.

While Skean et al. examined similarity across model architectures and modalities from an information processing viewpoint, our work concretely links representation geometry metrics to downstream model behavior and provides a foray into potential mechanisms underlying the representation geometry changes during pretraining and post-training stages. Thus, our findings offer a complementary view of how the training strategies lead to representational changes in LLMs and how those lead to emergent behaviors.

Intervention studies

We thank the reviewer for this insightful suggestion and pointing to this work by Arefin et al. We agree with the reviewer that our findings on multiple phases could lead to spectral interventions and regularization strategies for accelerating training and improving LLM generalization. We hypothesize that some of this regularization might be implicit in current objective functions, such as the equivalence between the DPO objective and contrastive loss that reveals an implicit entropy-seeking prior. The interplay between these implicit biases and explicit regularization warrants further study, which we believe is beyond the scope of this paper and leave for future work. However, we agree with the reviewer that mentioning this potential impact of our results will strengthen our paper, and we will add this point and a citation to Arefin et al. in the discussion section.

We thank the reviewer once again for their overall positive rating and are grateful for their helpful suggestions. We believe that incorporating our results from preliminary experiments altering the eigenspectrum and GPT-2 training in the appendix as well as discussing the recent works from Skean et al. and Arefin et al. will strengthen our paper and improve its impact. We hope that the reviewer finds our responses to be helpful, and are happy to clarify any further questions during the discussion phase.

References:

1. Skean, O., Arefin, M.R., Zhao, D., Patel, N.N., Naghiyev, J., LeCun, Y. and Shwartz-Ziv, R., 2025. Layer by Layer: Uncovering Hidden Representations in Language Models. In Forty-second International Conference on Machine Learning.

We thank the reviewer for their response and for engaging with us during the discussion phase. We appreciate the positive feedback and believe that the revisions made in response to the reviewer’s suggestions will enhance the clarity and overall quality of the paper. We are grateful for the reviewer’s time and constructive comments.

We thank the reviewer for their astute observations, encouraging comments and helpful suggestions, and we are glad that they found our work to be interesting and timely. Below, we address all the concerns and questions raised by the reviewer.

Statistical analysis

We agree that adding statistical testing will strengthen our results. Below, we present spearman correlation analysis results for the following results:

1. Correlation with task performance: We compute the correlation between the spectral metrics and downstream task performance (SciQ and TriviaQA) across different pretraining checkpoints for Pythia models.

SciQ

TriviaQA

2. Correlation with pass@256 performance: Since both RankMe and pass@256 values are nearly monotonically decreasing, the spearman rank correlation is ~1. We will extend our RLVR analysis to more checkpoints and compute the correlation for these results, and update Fig 5C.

We propose to add these results to the Appendix, in order to strengthen our claims about the correlational behavior of our findings.

Experiment details for computing spectral metrics

We apologize for the lack of clarity in these experimental details. For all experiments, we used 15k sequences from the Fineweb dataset and extracted the last layer representations after processing sequences of 512 tokens. We will add these details in Section 2.1 and amend it to have references to the ablation experiments in Fig 7 and Fig 8. We hope that this also clarifies the related concerns about the results in Section 3.4.

Clarification about grokking

While we agree the mechanism of grokking here differs from those studied in models of modular addition, we would like to clarify that TriviaQA is a factual question-answering dataset, not multiple-choice. Consequently, TriviaQA performance is measured by Exact Match metric via LLM harness. Nevertheless, we concur that "grokking" is an overloaded term and used it to describe a sudden increase in 0-shot task performance. To improve clarity and avoid overstatement, we will replace "grokking" with "rapid increase in 0-shot performance" in the text.

Clarification about Fig 1

Fig 1C summarizes our key findings graphically without diving into the technical details, using "representation complexity" to denote representation space entropy and only distinguishing between pre-training and post-training stages on the x-axis. While we believe it offers a high-level overview, we acknowledge the reviewer's feedback regarding clarity and avoiding speculation. Therefore, we will clearly differentiate between established and speculative claims in the figure and note this distinction in the caption.

Clarification about Section 3.2 results

1. Infini-gram results: We apologize for the lack of clarity in connecting our results to the claim about short-sequence memorization. The infini-gram model, despite its name, does not use a very large context window for predictions. It is a generalization of the n-gram model, where 'n' is not fixed and instead set to the length of the longest suffix found in the corpus. For samples from the TriviaQA dataset, we analyzed the suffix length from the question used to predict tokens in the answer and found that the mean was 4.21 tokens (median 4.0, max 12). A more detailed breakdown of the frequency of different suffix lengths is presented below. | Suffix Length | Frequency (%) | | :--- | :---: | | $\le$3 | 25.41 | | 4 | 46.61 | | 5 | 16.71 | | 6 | 6.08 | | 7 | 2.40 | | 8 | 1.46 | | $>$8 | 1.34 |

This analysis indicates that the infini-gram model's predictions primarily depend on short- to mid-range dependencies. While we acknowledge the reviewer's valuable suggestion of a more nuanced analysis using bigrams and trigrams, the computational burden of training several n-grams on the entire pretraining dataset is not feasible within the given rebuttal timeframe. Therefore, we have restricted our analysis to using infini-grams to calculate the memorization metric, and we will clarify the precise meaning of this measure in the updated manuscript.

2. Context length dependencies for SciQ and TriviaQA: The reviewer is correct that a key difference between SciQ and TriviaQA is that SciQ is a multiple-choice task, whereas TriviaQA is not. However, TriviaQA is included in the LongBench dataset [1] and is considered to require long-context understanding. In contrast, Lou et al. showed that SciQ accuracy, unlike TriviaQA performance, remains relatively stable even when using sliding window attention (Tables 2 & 4 in [2]), indicating that long-context information is not essential. Based on these results, we believe the performance discrepancy between the two benchmarks is in part due to the different dependency lengths needed to answer questions in their respective datasets. But, there may also be an effect due to the increased reasoning requirement to answer questions in the TriviaQA dataset, and we will note this as another possible interpretation in the updated text.

We hope that this clarification about the context-length required for the two tasks and the relationship between the memorization metric and short-context understanding ability will help resolve the logical jump pointed out by the reviewer.

Control experiment for clarifying the role of cross-entropy

We agree that using an objective other than cross-entropy, e.g. mean squared error (MSE), would be a natural control for our claim. We extended our theoretical analysis in Appendix B to derive the learning dynamics under an MSE loss. Simulations in the analytically tractable settings show that the eigenvalues monotonically increase (with $\sigma_1$ increasing faster than $\sigma_2$ initially) and eventually saturate. Consequently, the RankMe first decreases (warmup phase), then increases (entropy-seeking phase), and finally saturates. Notably, unlike the cross-entropy case, there is no smoothness-seeking phase when using MSE. This result further supports our claim that cross-entropy loss's learning dynamics are a key factor in the multiphase changes in representation geometry. We will add the MSE learning dynamics derivation and plots to the Appendix.

Clarification about Section 3.4 results

1. Ood nature of AH and AF: As mentioned in Section A.2, AH is a chat dataset aimed at incorporating helpfulness and harmlessness, while AF is a generic chat dataset. The out-of-distribution (ood) nature spawns from the fact that the chat responses in the dataset are qualitatively different. While responses in AF are more geared towards short and concise responses, responses in AH are more conversational, clarifying and thought-provoking. Furthermore, the prompts in AF are relatively straightforward and atomic in nature, i.e. they ask the model to do a single task, whereas the prompts in AH are complex user queries and often require deeper reasoning and understanding of the underlying intent. These qualitative differences make the two datasets distinct and inherently ood to each other.

2. Intuitive explanation of the chat winrate results: The winrate reported in Fig 5B (bottom) is computed by comparing the AH fine-tuned and AF fine-tuned model responses for prompts from a third, novel dataset: AlpacaEval. This comparison is done with a larger LLM judge which picks the more appropriate answer between the two. It is worth noting that AlpacaEval, which has prompts of similar style and complexity to AF, is more out-of-domain for AH fine-tuned models than AF fine-tuned models. This is reflected in the low winrates (6-14%) of AH fine-tuned models with respect to AF fine-tuned models. Our key observation is the decrease in this winrate with more pretraining of the base model, suggesting that ``overtrained'' models (those with higher RankMe values) perform better on in-domain tasks but worse at generalizing to out-of-domain data.

3. Intuitive explanation of the DPO objective reformulation: Equation 5 establishes the equivalence between the DPO objective and the contrastive learning objective. This equivalence suggests similar learning dynamics for both, which could explain the empirically observed entropy-seeking phase during the DPO training stage. We aim to further elaborate on this relationship and connect DPO learning dynamics to previous works on self-supervised learning dynamics in the Appendix of our paper's final version.

We thank the reviewer once again for their astute comments and are extremely grateful for their helpful suggestions. We believe that adding the statistical testing results and clarifications will strengthen our work and improve the clarity of the presentation. We hope that the reviewer finds our responses to be helpful, and are happy to clarify any further questions during the discussion phase.

References:

1. Bai, Y., Lv, X., Zhang, J., Lyu, H., Tang, J., Huang, Z., Du, Z., Liu, X., Zeng, A., Hou, L. and Dong, Y., 2023. Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint arXiv:2308.14508.
2. Lou, C., Jia, Z., Zheng, Z. and Tu, K., 2024. Sparser is faster and less is more: Efficient sparse attention for long-range transformers. arXiv preprint arXiv:2406.16747.

We thank the reviewer for their insightful comments and suggestions, and we are glad that they found our post-training results interesting and our work to be clear and timely. Below, we address all the concerns and questions raised by the reviewer.

Rogue dimensions

The reviewer is correct in noting that rogue dimensions in the LLM representation space affect the RankMe measure, leading to a lower effective rank than the representation dimensionality. This is particularly pronounced in models in the final stages of pretraining, which exhibit a much lower RankMe (computed using the entire eigenspectrum) and therefore indicate an inefficient use of their representation capacity, consistent with Timkey & van Schijndel's findings. In line with the reviewer’s comment, we will add an Appendix figure illustrating the impact of removing these rogue dimensions on RankMe.

In contrast, $\alpha$-ReQ quantifies the decay rate of the eigenspectrum tail, specifically the slope of the power-law decay fit to the eigenspectrum after the top 10 eigenvalues. By disregarding the top (rogue) eigendirections, $\alpha$-ReQ is less influenced by a few rogue directions. However, a strong inverse correlation between the two metrics indicates that the LLM representation space eigenspectrum is significantly skewed beyond just the top rogue dimensions.

Our work demonstrates that analyzing the degree of this skewness using our proposed spectral metrics, and observing its evolution across pretraining and post-training, provides deeper insights into model behavior. We concur with the reviewer that integrating these distinctions between the two metrics and their connection to Timkey & van Schijndel's work into the discussions will enhance our paper, and we propose to implement these changes in the final version.

$\alpha$-ReQ vs RankMe

While $\alpha$-ReQ is more robust to outlier rogue dimensions than RankMe and is based on an empirically observed power-law distribution of tail eigenvalues, both metrics are strongly anti-correlated and largely interchangeable. However, it is worth noting that RankMe is more sensitive to the ambient dimensionality of the representation space and dataset variations (Fig. 8), making cross-model comparisons challenging (Fig. 9).

Fig 4D clarification

We sincerely apologize for this confusion. Figure 4 presents the simplified setting that seeks to verify our analytical results about learning dynamics of cross-entropy loss. In these particular plots (and in Fig 10a), we set the feature dimensionality, $d=2$. Therefore, the eigenspectrum has only 2 eigenvalues, $\sigma_1, \sigma_2$ $\textemdash$ the ones that are plotted in Fig 4D. We will clarify these points in the updated manuscript.

Fig 5B clarification

We apologize for the confusion in these results. The reviewer has correctly noted that we fine-tuned (SFT) models on two different datasets: Anthropic-HH (AH) and AlpacaFarm (AF). Both these sets of models were then evaluated on a third, out-of-distribution (novel) dataset: AlpacaEval.

Moreover, it is worth noting that AlpacaEval is more out-of-domain for AH models than for AF models. Anthropic-HH is a chat dataset focused on helpfulness and harmlessness, with conversational, clarifying, and thought-provoking responses to complex user queries. In contrast, AlpacaFarm and AlpacaEval are generic chat datasets with straightforward, atomic prompts that elicit short, concise responses. The qualitative differences in chat responses and prompt complexity make AlpacaEval a more significant shift for AH models.

The winrate is computed by comparing the AH and AF model responses for the AlpacaEval prompts. This comparison is done with a larger LLM judge which picks the more appropriate answer among the two. As described above, AlpacaEval is more out-of-domain for AH models than for AF models. This is reflected in the low winrates (6-14%) of AH fine-tuned models with respect to AF fine-tuned models. Our key observation is the decrease in this winrate with more pretraining of the base model, suggesting that “overtrained” models (those with higher RankMe values) perform better on in-domain tasks but worse at generalizing to out-of-domain data. We will incorporate these clarifications about the experiment design in the final version of our paper.

We thank the reviewer once again and are grateful for their helpful suggestions. We believe that incorporating discussion of rogue dimensions and the nuances about the two spectral metrics will strengthen our paper and improve its readability. We hope that the reviewer finds our responses to be helpful, and are happy to clarify any further questions during the discussion phase.