AI Papers Reader

Personalized digests of latest AI research

View on GitHub

LLMs: The Geometry of Tokens Reveals Insights into Next-Token Prediction

A new paper from researchers at the University of Amsterdam and Area Science Park in Trieste, Italy, explores the hidden geometric relationships within large language models (LLMs). Their findings reveal a surprising connection between the geometric properties of token embeddings within the model and the model’s ability to accurately predict the next word in a sequence.

The study, titled “The Geometry of Tokens in Internal Representations of Large Language Models,” uses a novel approach to investigate LLMs. Instead of focusing solely on the output of the model, the researchers delve into the internal layers, analyzing how the representations (embeddings) of individual tokens evolve as they move through the network. They view the model as a dynamical system, where tokens are like particles interacting and evolving over many layers. The distribution of these token embeddings is called the “empirical measure,” and its evolution through the network is crucial for understanding prediction accuracy.

The team utilized three key geometric metrics to probe the empirical measure:

  1. Intrinsic Dimension (ID): This measures the effective dimensionality of the space where the token embeddings reside. A higher ID suggests more complex relationships between the tokens.
  2. Neighborhood Overlap (NO): This quantifies the consistency of token relationships across layers. High NO suggests that relationships learned in earlier layers persist and are refined in later layers.
  3. Cosine Similarity: This standard metric assesses the similarity between token vectors, revealing patterns of alignment or orthogonality.

To test their approach, the researchers systematically shuffled the tokens within input sentences, disrupting the syntactic and semantic structure. This allowed them to observe how the geometric metrics changed when the model processed semantically meaningless input.

Concrete Examples:

Imagine the sentence “The quick brown fox jumps over the lazy dog.” The researchers embed each word (token) into a high-dimensional vector space. In early layers, the vectors for related words might be clustered together, and the ID would be relatively low, reflecting the coherent structure of the sentence. However, after passing through multiple layers, the vectors for “quick” and “lazy” may now be farther apart (higher ID) while those for semantically related words like “quick” and “jumps” may be closer together. This change reflects the model’s understanding of semantic relationships and syntactic structure. In contrast, if the sentence were shuffled into “dog over jumps fox brown quick lazy the”, the geometric structure would be drastically altered, with less clustering and a higher ID.

Their findings showed several key correlations:

  • ID peaks in the middle layers: The ID of token representations generally exhibited a peak in the middle layers, suggesting that the representation space expands to capture intricate semantic relationships before collapsing in later layers due to “rank collapse”, an often-observed phenomenon in transformers. This peak was accentuated with increased token shuffling.
  • Loss correlates with ID: A statistically significant correlation exists between the average cross-entropy loss (a measure of prediction uncertainty) and the ID of token representations, especially around the peak layer. Prompts resulting in higher prediction uncertainty had tokens represented in higher-dimensional spaces. This indicates that higher uncertainty prompts require a more complex representation space.
  • NO shows layer-consistency for structured data: Neighborhood overlap was highly consistent for the unshuffled sentences, indicating that structured relationships between tokens were maintained across layers. Shuffling disrupted this consistency, reflecting the model’s reliance on syntactic and semantic structure.

Implications:

This work suggests that the geometric structure of internal token representations in LLMs is deeply linked to the models’ predictive capabilities. By analyzing these geometric properties, researchers can gain insights into how LLMs process information and build meaning, leading to a better understanding of the inherent dynamics within these complex systems. Further research is needed to explore the implications of these findings for model development and interpretability. The study highlights a promising new avenue for mechanistic interpretability of LLMs.