Probability Signature: Bridging Data Semantics and Embedding Structure in Language Models
This work provides a novel understanding of embedding formation mechanisms, which is foundational for improving interpretability in AI, though it is incremental in building on existing theories of semantic relationships.
The paper tackled the problem of understanding how embedding structures in language models relate to semantic relationships by proposing probability signatures derived from data distribution. The results showed that these signatures significantly influence embedding structures, with experiments on composite addition tasks and large language models demonstrating faithful alignment, particularly in capturing pairwise similarities.
The embedding space of language models is widely believed to capture the semantic relationships; for instance, embeddings of digits often exhibit an ordered structure that corresponds to their natural sequence. However, the mechanisms driving the formation of such structures remain poorly understood. In this work, we interpret the embedding structures via the data distribution. We propose a set of probability signatures that reflect the semantic relationships among tokens. Through experiments on the composite addition tasks using the linear model and feedforward network, combined with theoretical analysis of gradient flow dynamics, we reveal that these probability signatures significantly influence the embedding structures. We further generalize our analysis to large language models (LLMs) by training the Qwen2.5 architecture on the subsets of the Pile corpus. Our results show that the probability signatures are faithfully aligned with the embedding structures, particularly in capturing strong pairwise similarities among embeddings. Our work uncovers the mechanism of how data distribution guides the formation of embedding structures, establishing a novel understanding of the relationship between embedding organization and semantic patterns.