Polar probe linearly decodes semantic structures from LLMs

Technical Summary: Polar Probe Linearly Decodes Semantic Structures from LLMs

Problem Formulation
The paper addresses the fundamental question of how Large Language Models (LLMs) bind concepts to form complex semantic structures. While previous work has demonstrated that syntactic structures (e.g., subject-verb dependencies) can be linearly decoded from LLM activations using geometric principles (Hewitt & Manning, 2019; Diego-Simon et al., 2024), it remains unclear if similar principles apply to semantic binding. The authors hypothesize that LLMs represent the existence and type of relations between entities not through complex, non-linear mechanisms, but via a simple geometric code: the Euclidean distance between entity embeddings encodes the existence of a relation, while the relative direction encodes the relation type.

Methodology
To test this hypothesis, the authors introduce a Polar Probe, a linear transformation trained to map LLM layer activations into a probe space where the relational graph can be recovered.

1. Datasets: The study utilizes a synthetic dataset spanning five distinct semantic domains:

   • Euclidean domains: Variable ordinality (mathematical variables), spatial layouts (2D object arrangements), and thematic roles (social interactions).
   • Non-Euclidean domains: Family trees (non-commutative, many-to-many relations) and metro maps (network-based shortest paths).
   • Each sample consists of a natural language description of a relational graph. To prevent the model from learning surface-level correlations, the order of relations and entities is randomized, and multiple textual descriptions are generated for the same graph.

2. Probe Architecture:

   • Structural Loss ($L_s$): Optimizes the Spearman rank correlation between the predicted pairwise distance matrix and the ground-truth shortest-path distance matrix of the semantic graph.
   • Angular Loss ($L_a$): Optimizes the alignment between the difference vector of two entity embeddings ($\delta_{ij}$) and a learned prototype vector ($p_r$) specific to the relation type $r$.
   • The probe is trained independently on each layer of various LLMs (Llama3-8B, OLMo-7B, and the Pythia family).

3. Evaluation:

   • Decoding Performance: Measured via Spearman's $\rho$ for relation existence (distance ordering) and relation type (direction alignment).
   • Generalization: Tested on Out-of-Distribution (OOD) entity names and relation surface forms.
   • Causal Intervention: The authors intervene on the hidden states by steering activations along the learned prototype directions to observe shifts in downstream Question-Answering (QA) performance.
   • Correlation: Probe errors are correlated with the logit probability of correct answers in QA tasks.

Key Results

• Linear Decodability: The true semantic structures can be linearly recovered from LLM activations. The Polar Probe successfully decodes both relation existence and type across all five domains.
• Layer Emergence: The geometric code is most prominent in the middle layers of pretrained LLMs (peaking around layers 12–15 for Llama3-8B). In randomly initialized models, decoding performance remains at chance levels, indicating the code emerges during pretraining.
• Model Size and Training: Performance scales with model size (from 14M to 6.9B parameters) and increases with pretraining steps, suggesting the polar code is a learned capability dependent on model capacity and training duration.
• Generalization: The probes generalize to new entities and relation surface forms, though performance degrades modestly with OOD entities and more significantly with OOD relation types. Performance also declines as the number of entities in the graph increases, likely due to the expanding combinatorial space.
• Causal Influence: Steering the model's activations along the learned prototype directions causally shifts the probability of correct answers in QA tasks, confirming that these geometric representations are functionally relevant to the model's reasoning.
• Subspace Disjointness: The learned subspaces for different semantic domains are largely disjoint in the activation space, with the notable exception of spatial layouts and variable ordinality, which show significant overlap (supporting the superposition hypothesis).

Significance and Claims
The paper claims to demonstrate that LLMs learn to build complex semantic structures by binding representations through a simple geometric principle. Specifically, it extends the scope of linear probing beyond syntax to a broader class of directed, labeled graphs.

• Beyond Syntax: Unlike prior work limited to syntactic trees, this approach successfully decodes semantic structures that are not necessarily tree-structured and include non-Euclidean topologies (e.g., family trees, metro maps), although with some limitations in accuracy for non-commutative relations.
• Beyond Knowledge Retrieval: The study targets compositional structures that could not be memorized during training (due to randomized entities and relations), providing evidence for how models perform in-context learning and binding.
• Symbolic-Connectionist Tension: The findings suggest a resolution to the tension between symbolic and connectionist approaches, proposing that symbolic structures can be represented within vectorial systems via linear subspaces and geometric relationships.

The authors remain modest regarding the limitations of the method, noting that the Polar Probe relies on linear distances and is tailored for Euclidean graphs. Consequently, it cannot perfectly represent entire non-commutative, many-to-many relations or complex shortest-path distances found in non-Euclidean domains, though it captures local coherence within them.