The Geometry of Tokens in Internal Representations of Large Language Models

2025-01-17

AI Summary

This study investigates how the geometric structure of token embeddings in large language models (LLMs) influences next-token prediction performance. Method: We propose a systematic empirical measurement framework that integrates geometric metrics—including intrinsic dimensionality, neighborhood overlap, and cosine similarity—and quantitatively correlates them with per-layer cross-entropy loss for the first time. Contribution/Results: We find that high prediction loss consistently corresponds to higher-dimensional embedding spaces, revealing an implicit link between geometric dimensionality and model uncertainty; this correlation vanishes under token shuffling, confirming the decisive role of syntactic and semantic structure in shaping embedding geometry. Furthermore, layer-wise geometric evolution analysis characterizes a dynamic transition from discretized to manifold-like representations across layers. These findings establish embedding geometry as an interpretable proxy for predictive performance, offering a novel perspective for model diagnosis and optimization.

Technology Category

• Natural Language Processing: (Large) Language Models
• Machine Learning: Large Multimodal Models (LMMs)
• Knowledge Representation and Reasoning: Geometric, Spatial, and Temporal Reasoning

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

Abstract

We investigate the relationship between the geometry of token embeddings and their role in the next token prediction within transformer models. An important aspect of this connection uses the notion of empirical measure, which encodes the distribution of token point clouds across transformer layers and drives the evolution of token representations in the mean-field interacting picture. We use metrics such as intrinsic dimension, neighborhood overlap, and cosine similarity to observationally probe these empirical measures across layers. To validate our approach, we compare these metrics to a dataset where the tokens are shuffled, which disrupts the syntactic and semantic structure. Our findings reveal a correlation between the geometric properties of token embeddings and the cross-entropy loss of next token predictions, implying that prompts with higher loss values have tokens represented in higher-dimensional spaces.

Problem

• Permutation Effect
• Token Prediction
• Representation Evolution

Innovation

• Transformer Models
• Token Representation Evolution
• Prediction Error Correlation