Shared Global and Local Geometry of Language Model Embeddings

Thank you for your response. Everything looks good to me. I believe this is a strong paper and I will maintain my current score.

Q4: For the proof in Appendix A, can the authors explain line 407 ?

Adding $\epsilon I$ is a common trick (ridge regularization, aka Tikhonov regularization), which provides a few benefits but here are two:

1. By defining $C_i = Z_i^TZ_i + \epsilon I$, every eigenvalue of $Z_i^TZ_i$ gets shifted by $\epsilon$, and thus $C_i$ becomes strictly positive definite (becoming invertible), and allowing a unique solution.
2. In line 410, we are interested in minimizing $w^TCw = w^TZ^TZw = ||Zw||^2$. By defining $C = Z^TZ + \epsilon I$, this becomes $||w^TCw||^2 = ||Zw||^2 + \epsilon ||w||^2$, where the first term $||Zw||^2$ is our reconstruction error and the $\epsilon ||w||^2$ term is a ridge penalty term, where the ridge penalty term similarly allows for the optimization objective to have a unique solution.

In practice, the use of a pseudo-inverse or SVD is also acceptable and would likely lead to similar results.

Lastly, note that LLE (as well as our closed form approach) is an approximate reconstruction of token embeddings to begin with.

Q6: With regards to the steering vector, do the authors have any results applying the transfer on models from different families ? i.e. llama -> gpt ?

In this work we did not study cross-family transferability or geometric similarities. This was mostly because our approach requires matching tokens between embedding spaces, and thus requires LMs that share the same tokenizer (ie, same model family).

Apologies for the late response. I had to re-evaluate this paper based on some of the responses.

While I am happy with the responses to my other questions, I would like to respectfully request that the authors revisit Question 5, which I believe was not adequately addressed in the current draft. After a closer examination of Section 4, I am concerned that the primary findings on “Model Families Share Similar Token Orientations (Usually)” may largely be a byproduct of shared training data across models, rather than an intrinsic or architecture-driven phenomenon. This shared global geometry regardless of model is discussed in the aforementioned works.

In particular, recent work such as Wu et al. (2024), which explores Neural Collapse in language models, offers a data-driven perspective on embedding geometry. Their framing of next-token prediction as a classification problem — and the finding that last-layer features collapse for contexts sharing the same next token — provides a strong precedent for viewing embedding geometry as emergent from dataset structure. Additionally, Zhao et al. (2024) offers both theoretical and empirical evidence that the geometry of unembedding vectors (referred to in their work as last-layer weights) is tightly linked to the distribution of next-token support sets in the training corpus.

While the current paper presents a compelling empirical study, it does not engage with these relevant lines of work. The conceptual overlaps are significant: both Wu et al. and Zhao et al. propose that token representations are shaped by token-level supervision signals and co-occurrence patterns, which is fundamentally the same mechanism that could explain the shared global geometry reported in this paper (different models sizes of the same family most likely had the same training data). In fact, the phenomenon of aligned token embeddings across models is explicitly theorized and empirically analyzed in Zhao et al. (2024) through a optimization framing — yet the authors do not cite or discuss them. I believe a more thorough engagement with this literature is essential to properly contextualize and substantiate the paper’s claims.

Given the conceptual proximity between the current work and the recent studies cited above, I find this omission troubling. I understand that the literature in this area is vast and fast-evolving, and it is reasonable that some relevant work may be unintentionally overlooked. However,I am particularly troubled by the fact that this issue was explicitly raised in my original review as Question 5, yet it was not addressed in the author response at all. Given how directly related these prior works are to the claims made in this paper, I believe a thoughtful engagement with them is necessary. As it stands, given that this global geometry of LLMs is discussed in other forms in previous work, it takes away from some of the novelty of the paper regarding this particular topic. One of the reasons that I assigned a high score to this work was because I was unsure of the similarity between the paper and previous literature and still found the idea to be novel. Given the lack of discussion, I am reducing my score to a 6.

Thank you for your response, and for the two pointers. Both papers were great!

This is not true for the unembedding given that for a token $t_j$, the unembedding vector $h^{L-1}$...

We first wish to quickly clarify a minor point regarding your response above regarding “unembedding vector $h^{L-1}$” – although it’s possible that we’re just confusing ourselves with notations, but hopefully this also clarifies the notations used in the rest of our response below.

After going through $L$ transformer blocks, LMs build a context representation in its penultimate layer ($\mathbf{h}^{L-1} \in \mathcal{R}^d$). This context representation is mapped back to a “vocabulary space” using unembedding vectors, $\mathbf{U} \in \mathcal{R}^{V \times d}$. The softmax on the result ($\mathbf{Uh}^{L-1} \in \mathcal{R}^V$) predicts the next token. So, while the argument to the softmax ($\mathbf{Uh}^{L-1}$) is context dependent, the unembedding weights $\mathbf{U}$ themselves are not.

Now indeed, there is an interesting, symmetric relationship between context representations ($\mathbf{h}^{L-1}$) and the unembedding layer ($\mathbf{U}$). As Zhao et al. puts it, $\mathbf{h}^{L-1}$ and $\mathbf{U}$ can be thought of as a log-bilinear model. We provide a similar characterization in Equation 6. Perhaps this symmetric, bilinear relationship is what you mean by unembedding vectors being context dependent.

Indeed, the two works you point to study the geometry of either the penultimate layer $\mathbf{h}^{L-1}$ or the bilinear model $\mathbf{Uh}^{L-1}$. First, regarding neural collapse, it is neat to see neural collapse studied in language models! There can definitely be an interesting relationship between neural collapse and geometry of unembeddings. Namely, if we put aside for a moment the strict criteria for neural collapse (ex: simplex ETF) and simply think of neural collapse as an “optimal representation” that can be learned in $\mathbf{h}^{L-1}$, then it is certainly expected that models that reach neural collapse will also have similar geometry in their symmetric counterpart, $\mathbf{U}$. It would be interesting to empirically confirm this – that models with stronger signatures of neural collapse will also have more similarities in the geometric signatures that we study.

Regarding Zhao et al., there are definitely strong connections between their theoretical findings and our work. For instance, their finding that $\mathbf{Uh}^{L-1}$ will be dominated by a low-rank component, driven by word co-occurrences, could be a potential explanation for the similarity of geometries seen in our work. Finding further connections between the two works could be a great contribution.

Yes, there is a close relationship between the geometry of context representations ($\mathbf{h}^{L-1}$) and unembeddings ($\mathbf{U}$). Yes, all 3 papers study the similarities in the geometry of either $\mathbf{h}^{L-1}$, $\mathbf{Uh}^{L-1}$, or $\mathbf{U}$.

However, some specific contributions in our work include:

• Simple ways of characterizing and comparing global and local geometry (ex: LLE)
• Relationship between low intrinsic dimensions and semantically coherent clusters
• Transferability of steering vectors

Lastly, regarding omitting a response to question 5 - genuinely, there was no malicious intent. We did not understand what was meant by “[authors] claim that token embeddings are independent of contextual usage” and how it connects to the two papers. Again, in discussing why steering vectors are transferable (Section 6, Equation 6), we discuss the bilinear relationship between $\mathbf{h}^{L-1}$ and $\mathbf{U}$, just like Zhao et al.’s log-bilinear setup, and even use this argument in our response to reviewer tRQC for why untied models have dissimilar embeddings but similar unembeddings.

(Realizing as we write this up now - we can see how our response to reviewer tRQC regarding the bilinear relationship in $\mathbf{h}$ and $\mathbf{U}$ combined with the omission to question 5 could have looked suspicious - maybe this is what you are referring to.)

For what it’s worth, we are used to rebuttals having character limits (ex: 6,000 characters for Neurips), which historically meant that one could not always respond to every point and had to prioritize what to respond to. While COLM does not have a character limit, the instructions specifically say “Regarding length: we strongly recommend against writing long responses. Please be succinct and to-the-point in your response. This lowers the barrier for reviewers to engage meaningfully.” While COLM is new and thus does not have established norms yet, we took this to mean that we should follow the norms of other ML conferences.

All of this is to say, we genuinely did not mean harm and apologize for the misunderstanding. We kindly ask that your score reflects what you view as our contributions (in light of all relevant prior work), and not any misunderstandings.

Thank you!

I am happy to elabotrate. If I have the same token, in different sequences, the embedding vector for that token is still the same I assume. this is what I was referring to as the context independance of token embeddings.

This is not true for the unembedding given that for a token token $t_j$, the unembedding vector $h_{j}^{L-1}$ depends on the tokens appearing before it in the sequence $[t_0, t_2, \dots , t_{j-3}, t_{j-2}, t_{j-1} ]$.

Thank you for the detailed responses and clarification, and apologies for my earlier misunderstanding. Most of my previous concerns have now been addressed. I also appreciate the authors’ willingness to incorporate more high-level explanations in the revision. Given these clarifications and the authors' intent to improve the presentation and discussion in the revision, I have decided to revise my rating from 4 to 7.