Less is More: Local Intrinsic Dimensions of Contextual Language Models

We thank Reviewer ​​oyKS for their feedback and for highlighting the potential of our work to advance the study of geometric representations in language models. We are particularly grateful for the appreciation of our effort to connect intrinsic dimensionality with diverse learning dynamics such as fine-tuning, grokking, and overfitting, and for acknowledging the value of studying the representation geometry of LLMs. Below, we address each concern in detail.

Addressing the Weaknesses

(1) Limited number of experiments per phenomenon

A broader empirical basis per phenomenon would certainly strengthen the generality of our claims. While we focused the main text on one representative setup per phenomenon for clarity and space, we do have additional results that support the generality of our findings.

For example, in the context of overfitting, we conducted further fine-tuning experiments on language modeling, where continued training past optimal validation performance led to an increase in local dimension after the initial drop. This further confirms that a dimension increase is a useful indicator of overfitting. We will include these results in the appendix of the revised version.

(2) Clarification of dataset and model choices

Thank you for requesting more transparency here. We will clarify these design decisions in the revision. Briefly:

• Our dataset selection was guided by the goal of interpreting training and fine-tuning dynamics in controlled, well-understood domains. SDG and MultiWOZ remain as the largest openly available task-oriented dialogue datasets, and they come with rich annotations (for example, EmoWOZ as an augmented version of MultiWOZ with additional emotion labels per user utterance). Dialogue provides a test-bed for many natural language processing phenomena, allowing us to study discourse-related challenges like natural language understanding, state tracking, and emotion classification in a realistic yet compact setting.

• The fine-grained and high-quality annotation of MultiWOZ/EmoWOZ allows us to train models with different goals on the same underlying dataset. We decided to cover sequence tagging and tracking via dialogue state tracking in the Trippy-R model, and sequence classification via emotion classification in the ERToD model. This, in turn, allowed us to isolate specific learning signals in fine-tuning and overfitting settings.

• For models, we selected representative examples of transformer-based encoder-only (BERT/RoBERTa) and a decoder-only model (GPT-2) to capture both sides of the architectural spectrum. These models are well-established and allow for interpretability across training regimes, since, in particular, RoBERTa-style models [a] are still widely applied as encoder models.

• To support our claims of wider applicability, we are currently running additional experiments on more recent models. Preliminary results on Phi-3.5-mini-instruct [b], a decoder model with 3.82B parameters and a hidden dimension of 3072, can be found in the table below. These results demonstrate that the TwoNN estimator, applied locally within our subsampling method, yields meaningful estimates even for latent spaces whose ambient dimension numbers in the thousands, a typical setting for modern LLM architectures.

In the case of the Phi-3.5-mini-instruct model, we implement the fine-tuning via LoRA with rank constraint r=16 on a portion of the MultiWOZ and Reddit training datasets, with other hyperparameters as described in the paper and taking a model checkpoint after 800 batches. During this fine-tuning, we observe the shifts in mean local intrinsic dimension on a subsample of the validation set of various datasets. In addition to the datasets in the paper, we present results for a self-collected dataset of ICLR 2024 titles and abstracts, for the sake of including a more recent dataset.

The following table contains the (Mean / Median / Std) of the local intrinsic dimension for base and fine-tuned (FT) versions of the Phi-3.5-mini-instruct model.

Note that our qualitative observations on the behaviour of the autoregressive model’s latent spaces under fine-tuning hold for this larger model as well: The dimension drops on the datasets that are related to the fine-tuning data (e.g., the dialogue datasets MultiWOZ and SGD become lower-dimensional in the latent space when tuning the model on MultiWOZ). At the same time, the dimension of unrelated datasets can be seen to increase or stay unchanged, similar to what happens when fine-tuning the notably smaller GPT-2 model. Moreover, we are currently working on an extension to even more recent models (e.g., from the Llama family), and we aim to include the respective results in our final manuscript.

(3) Appendix Figure 13 and fine-tuning generalization

The result in Figure 13 actually reinforces our core claim: it shows that the drop in local dimension primarily occurs within the domain of the fine-tuning data, while unrelated inputs remain largely unaffected or even increase in their mean local dimension. This is what we expect if the model adapts its representation geometry in a task-specific way. We will clarify this interpretation in the caption and appendix discussion to avoid confusion. In particular, we will make it clear that the discussion in Section 4.1 applies primarily to masked language models, and that decoder-only models appear to exhibit the phenomenon of increasing local dimension on non-training-related data.

(4) Discrepancy between Section 4.3 and 4.4, re: overfitting criteria

This is an important observation, and we appreciate the opportunity to clarify. In Section 4.3, we observe that validation loss increases while the non-differentiable downstream metric (Joint Goal Accuracy) still improves, indicating that when JGA is your target metric, validation loss is not a reliable stopping criterion in this case.

In contrast, in Section 4.4, the validation loss and accuracy (weighted F1) are tightly aligned; the respective maximum and minimum occur around the same training step. Hence, the loss serves as a valid stopping signal there, and the local dimension measures support this point. The difference in evaluation criteria is task-specific, and we will clarify this point more explicitly in the text.

(5) Interpretation of grokking as a phase transition

We agree that grokking is often interpreted as a phase transition. Our novel contribution here is to make this intuitive notion quantifiable: our local intrinsic dimension estimates provide a smooth, quantitative trace of the internal phase transition dynamics.

Crucially, our method enables this measurement using training data alone, and still anticipates the grokking event observed in validation accuracy. To the best of our knowledge, this is the first time such alignment can be observed at this granularity. We will revise the text to avoid overclaiming about prediction and emphasize the interpretability advantage.

Questions

(Q1) Choice of models and datasets

Please see our earlier explanation [see (2) above]. We will add a dedicated section explaining our design choices per setup in the revised paper.

(Q2) Suggested related work

We thank the reviewer for these additional references, and we will incorporate them into our Related Work section. While these works also explore intrinsic dimensions, they differ in key aspects: e.g., one focuses on sequence-level embeddings, while the other does not track dynamics during training.

Our method, in contrast, enables fine-grained, contextualized, local estimates over the course of optimization, which is essential for studying grokking or overfitting behavior. We are happy to discuss how our findings relate to and complement these efforts.

Limitations

Limitations and computational cost

Our approach is computationally intensive, which we discuss in our Limitations section, but at the same time, we provide a partial solution to the computational overhead via our sampling procedure. Moreover, we extensively analyse the stability of our estimates under sampling in our sensitivity analysis in Appendix A1.

The preliminary experiments with the Phi-3.5-mini-instruct model (3.82B parameters, 3072-dimensional latent space, for details see our rebuttal to Reviewer 67F6) demonstrate that this sampling approach applies to models with >1B parameters and embedding dimension in the thousands.

So while we largely agree with the reviewer’s assessment of the limitations, we believe that with minor modifications in the camera-ready version of the paper, these will be adequately addressed.

Conclusion

We thank the reviewer again for their time and engagement with our work. We hope that our clarifications and small additions to the revised manuscript address your concerns. Should these revisions meet your expectations, we would kindly ask you to consider increasing your score.

[a] Smarter, Better, Faster, Longer: A Modern Bidirectional Encoder for Fast, Memory Efficient, and Long Context Finetuning and Inference (arXiv:2412.13663)

[b] Phi-3 Technical Report (arXiv:2404.14219)

Thank you for the additional experiments. This has addressed my main concern, which is that the findings may not generalize. The addition of a 3B model is in particular good to see. With the additional promised changes, I am happy to raise my score to a 4.

We sincerely thank Reviewer rATb for the thoughtful and constructive review. We are grateful for the recognition of the strengths of our work, particularly your appreciation of the originality of our framing, the rigorous experimental design, and the clarity of the presentation. We are pleased the reviewer found practical utility in our method and robustness in our analysis. We would like to take this opportunity to address the specific points you raised:

Addressing the Weaknesses

Computational Overhead

We agree that computational cost is a potential barrier to large-scale adoption. However, we would like to emphasize that we already mitigate this issue through random subsampling of tokens, which significantly reduces overhead while maintaining stable estimates. Our sensitivity analysis (in Appendix A) empirically supports this approach, demonstrating that local dimension estimates are robust under sampling. We will further clarify this point in the revised manuscript to better guide practitioners on balancing accuracy and efficiency.

Relativity of Measurements

We acknowledge that absolute local dimension values depend on the choice of hyperparameters and model architecture. However, our extensive sensitivity analysis (Appendix A) shows that our findings are robust across a wide range of neighborhood sizes and subsampling rates. We will make this robustness more prominent in the main text.

Causality vs. Correlation

Regarding the distinction between causation and correlation in dimensionality reduction and generalization, while disentangling this relationship is beyond the scope of our current work, we agree that the question is critical. We will revise Section 4.6 to explicitly state this limitation and outline directions for future causal investigations. We are also happy to include a brief speculative discussion on possible mechanisms (e.g., information compression) by which reduced dimensionality might encourage better generalization.

Questions

Implications for Fine-Tuning / Regularization via LID:

This is an insightful question. While our current dimension estimation method is not differentiable (due to nearest-neighbor graph construction and the subsequent TwoNN estimator), we see this as an exciting opportunity for future work. Designing a differentiable surrogate or proxy loss to encourage local compression during training could open up new avenues for model regularization. We will expand our discussion of this direction in Section 5.

Heterogeneity of Local Dimensions Across Token Types

We have indeed explored the heterogeneity of local dimension distributions by grouping tokens according to POS tags. For example, we observed that on average, punctuation and stop words tend to occupy higher-dimensional regions, while content-bearing words such as nouns and verbs exhibit lower local dimensions. We will include a discussion of this phenomenon in the revised manuscript and add a supporting figure to the Appendix.

Connection to Parameter-Efficient Fine-Tuning (PEFT)

We appreciate this connection and have explored it preliminarily. The drop in mean local dimension on the fine-tuning set can be detected on LoRA-tuned models as well, in line with our observations in Section 4.1. Interestingly, we did not observe consistent differences in local dimensional dynamics between full and LoRA-based fine-tuning.

However, we agree that adapting the LoRA rank dynamically based on local geometry is a promising idea. We will add a note on this to the discussion section, citing the relevant work by Ed-dib et al. (2024), and marking it as an exciting direction for future research.

Encoder vs. Decoder Architectures:

While much of our analysis focuses on encoder models, we also included decoder-only models (Appendix C, Figures 12 and 13) and observed comparable dimensional drops on the fine-tuning set. We will clarify this in the main text and ensure that the decoder results in Appendix C.1 are better connected to the core narrative.

When comparing the distribution plots of a given dataset between the encoder-only RoBERTa and decoder-only GPT-2 model, one observes stark differences. For instance, the distribution of local estimates on the Reddit dataset in the GPT-2 latent space appears to be bimodal, which is less pronounced in the RoBERTa latent space. Investigating the reasons for these differences is an interesting direction for future research.

Additionally, we are currently working on an extension to even more recent decoder models (e.g., the Phi and LLaMA model family), and we aim to include the respective results in our final manuscript.

Limitations

Causal Limitation Not Explicitly Stated

We will explicitly note this limitation regarding causality between dimensionality and generalization in the revised limitations section.

We thank Reviewer rATb again for the detailed and supportive review.

We thank Reviewer PQs3 for their constructive feedback. We especially appreciate the recognition of our paper’s clarity, the novelty of the geometric perspective, and our effort to develop an unsupervised approach to interpreting LLM behavior via local intrinsic dimension (LID). We are encouraged that they found our main contributions clear and the methodological framing on understanding language models interesting. Below, we address each of the reviewer’s concerns in detail.

Addressing the Weaknesses

(1) Evaluation of modern LLM architectures

We would like to address the reviewer’s concern regarding the relevance to modern LLM architectures. We chose RoBERTa and its fine-tuned variants because our focus lies in understanding fine-tuning dynamics in contextual embedding spaces, settings where pretrained models are adapted to downstream tasks. These architectures remain widely used in practice, especially in encoder models [a] and in settings where compute or data constraints prohibit larger model training.

To demonstrate the feasibility of our approach on higher-dimensional latent spaces, we compute our local intrinsic dimension on the Phi-3.5-mini-instruct model [b], which has a hidden dimension of 3072.

The following table contains the (Mean / Median / Std) of the local intrinsic dimension for base and fine-tuned (FT) versions of the Phi-3.5-mini-instruct model. This demonstrates that our method provides reasonable estimates for higher-dimensional latent spaces produced from modern LLMs.

Moreover, we are currently working on an extension to even more recent models (e.g., Llama), and we aim to include the respective results in our final manuscript.

Our code is available as part of the supplemental material, and our local estimates computation is compatible with models from the HuggingFace transformers library. Thus, running our dimension estimation pipeline on larger models is mainly a matter of having compute available. According to Reviewers 67F6 and rATb, our experiments are comprehensive, and our model selection supports the conclusions drawn in our paper. Users of our method can always supply their own, larger model if applicable.

(2) Dataset selection and recency

We would like to address the reviewer’s concern regarding the selection of our datasets. Our goal was to capture contrasting behaviors in embedding space evolution, so we deliberately selected two complementary types of datasets:

(1) a dataset that was part of the original pretraining corpus (Wikipedia), and (2) a dataset that was released after the pretraining of the model concluded (Reddit) and hence not seen during training. This contrast allows us to probe how local intrinsic dimensions behave in both seen and unseen data regimes.

Our code is released as part of the supplemental material, and our local estimates computation pipeline is compatible with the HuggingFace datasets package. Thus, practitioners can easily extend the evaluation to their datasets. We demonstrate this above with our self-collected ICLR 2024 dataset containing recently released machine learning paper titles and abstracts.

(3) Computational cost of TwoNN estimation

We address the concern about computational cost in the limitations section. While TwoNN involves pairwise distance calculations, we emphasize that our method employs lightweight sampling strategies that drastically reduce computational demands. In practice, we use random token-level sampling, which makes the method feasible for large-scale experiments, as demonstrated across all tasks. Our new results on models from the Phi-family show that our method can be practically applied to modern LLM architectures.

(4) Validity of the TwoNN assumption in LLM embedding spaces

We interpret the value of the TwoNN estimator as a local dimension. That this interpretation hinges on certain assumptions is explicitly noted in the limitations section of the paper. However, the experimental results and their practical implications do not depend on this interpretation. This is why we view this limitation as relatively unimportant.

Still, we are happy to underpin our particular choice of local dimension estimator with theoretical findings:

• Locally uniform density of the embedding space: In Appendix B of “The geometry of hidden representations of large transformer models” (Valeriani et al., NeurIPS 2023), the authors demonstrate empirically that hidden states of Transformer models exhibit locally constant density, satisfying this condition.
• Poisson point process assumption: As noted in “The Geometry of Tokens in Internal Representations of Large Language Models” (Viswanthan et. al, arXiv:2501.10573), this condition is usually satisfied in our setup.

We will make these justifications explicit in the updated manuscript and add references and discussion.

Questions

(Q1) On the relationship between local dimension and token entropy

We thank the reviewer for raising this potential connection. Indeed, there may be a relationship between token-level entropy and local intrinsic dimension, and this is a promising direction for future work.

However, the recent studies on entropy and reinforcement learning (RLVR) cited ([1], [2]) focus on reward optimization and exploration behavior in instruction-tuned or RL-finetuned models, scenarios that differ fundamentally from our setting. These papers were also released after our submission deadline. We agree this connection is worth exploring and have added the general setting of RL-tuning of LLMs to our future work section.

(Q2) Causal link between LID and generalization/grokking

We agree that a clearer discussion of why reduced LID corresponds to generalization would improve the paper. Our working hypothesis is that high-dimensional neighborhoods indicate overparameterized, poorly localized behavior, while very low LID reflects overcompression or memorization. There exists a “sweet spot” in LID, where the representation captures the essential degrees of freedom required for solving the task. Our experiments show that this intermediate range correlates with peak generalization. We will make this intuition more explicit, but one should note that investigating causality is beyond the scope of this paper, which is also acknowledged by Reviewer rATb.

Conclusion

We believe that we have adequately addressed the concerns regarding evaluation scope, dataset choices, computational feasibility, theoretical grounding, and interpretation of our results. We are committed to incorporating all necessary clarifications and extensions in the final revision and appreciate the reviewer’s suggestions.

If you find that these responses and planned improvements sufficiently address your concerns, we would be very grateful if you would consider raising your score.

[a] Smarter, Better, Faster, Longer: A Modern Bidirectional Encoder for Fast, Memory Efficient, and Long Context Finetuning and Inference (arXiv:2412.13663)

[b] Phi-3 Technical Report (arXiv:2404.14219)

We sincerely thank Reviewer 67F6 for their thoughtful and constructive feedback. We are especially grateful for the positive assessment of the paper’s clarity, novelty, and scientific merit, as well as for highlighting the value of linking the geometry of embedding spaces to model behavior.

Statistical testing

Regarding the point about statistical testing, we emphasize that the results presented in the grokking, training dynamics, and overfitting sections (4.2, 4.3, 4.4) are averages over multiple training runs, accompanied by their respective standard deviations. The relatively small magnitude of these deviations compared to the observed changes supports our characterization of differences as “significant”. However, we agree that formal statistical tests would underline our claims. In the revised version of the paper, we will include appropriate statistical analyses to complement our main findings, particularly in the sections concerning grokking behavior, training, and overfitting dynamics.

Dimension measures for larger language models

While our experiments focus on RoBERTa, GPT-2 and a small transformer used for grokking studies, the geometric perspective and methodology we propose are broadly applicable across architectures.

To support this claim, we are currently running additional experiments on more recent models. Preliminary results on Phi-3.5-mini-instruct [b], a decoder model with 3.82B parameters and a hidden dimension of 3072, can be found in the table below. These results demonstrate that the TwoNN estimator, applied locally within our subsampling method, yields meaningful estimates even for latent spaces whose ambient dimension numbers in the thousands, a typical setting for modern LLM architectures. Note that, for example, the 70B parameter Llama family, such as "meta-llama/Llama-3.1-70B-Instruct" has a hidden dimension size of 8192, which is within reach of our methods.

In the case of the Phi-3.5-mini-instruct model, we implement the fine-tuning via LoRA with rank constraint r=16 on a portion of the MultiWOZ and Reddit training datasets, with other hyperparameters as described in the paper and taking a model checkpoint after 800 batches. During this fine-tuning, we observe the shifts in mean local intrinsic dimension on a subsample of the validation set of various datasets. In addition to the datasets in the paper, we present results for a self-collected dataset of ICLR 2024 titles and abstracts, for the sake of including a more recent dataset.

The following table contains the (Mean / Median / Std) of the local intrinsic dimension for base and fine-tuned (FT) versions of the Phi-3.5-mini-instruct model.

Note that our qualitative observations on the behaviour of the autoregressive model’s latent spaces under fine-tuning hold for this larger model as well: The dimension drops on the datasets that are related to the fine-tuning data (e.g., the dialogue datasets MultiWOZ and SGD become lower-dimensional in the latent space when tuning the model on MultiWOZ). At the same time, the dimension of unrelated datasets can be seen to increase or stay unchanged, similar to what happens when fine-tuning the notably smaller GPT-2 model.

These results and the discussion on larger models will be added in the camera-ready version of our paper. We are currently extending our analysis to more recent models (e.g., from the Llama family), and intend to include these results in the final version of our manuscript.

Our choice of RoBERTa was motivated by its continued use in research and applications where encoder-based models are needed [a] and autoregressive alternatives are less suitable, such as in information retrieval, which is an important component in Retrieval Augmented Generation. We will clarify this rationale and better situate the scope of our contributions in the revised manuscript.

Layout

Finally, we thank the Reviewer for the suggestion regarding figure layout and formatting. We will enlarge the visuals and adjust their placement to improve readability in the revised manuscript.

Once again, we are grateful for the constructive feedback and the encouraging evaluation of our work.

[a] Smarter, Better, Faster, Longer: A Modern Bidirectional Encoder for Fast, Memory Efficient, and Long Context Finetuning and Inference (arXiv:2412.13663)

[b] Phi-3 Technical Report (arXiv:2404.14219)