Divergent large language model predictions from convergent representations in ambiguous word pairs

How do large language models (LLMs) resolve lexical ambiguity? Specifically, do the geometric representations of ambiguous words (like homonyms and polysemes) in the final layers of a model accurately reflect the semantic distinctions the model uses to generate predictions, or is there a disconnect between internal geometry and behavioral output?

The researchers analyzed three decoder-only models of varying scales: GPT-2-Small (117M), Llama-3.2-3B, and Qwen2.5-32B. They tracked the internal representations of ambiguous word pairs across every layer using three primary methods:

• Logit Lens: Projecting intermediate hidden states into the vocabulary space to observe how predictions evolve.
• Activation Patching: Causal intervention where activations from one context are swapped into another to see if they force the model to adopt the source context's prediction.
• Single-layer Ablation: Zeroing out specific layers to measure how much each layer contributes to the model's ability to maintain disambiguation.

The study reveals a consistent pattern across all tested models: representations of ambiguous words become maximally distinct in the middle layers but then partially reconverge in the final layers. Crucially, the KL divergence between the model's next-token predictions—the behavioral manifestation of disambiguation—reaches its maximum in these same final layers. Activation patching confirms that these late-layer representations, despite appearing more similar in embedding space, are causally responsible for the model's final output. This suggests that semantic information is not lost in the late layers; rather, it becomes encoded in a way that is increasingly opaque to simple geometric metrics like cosine similarity.

This research challenges the common practice of using late-layer cosine similarity as a proxy for semantic meaning in tasks like semantic search, retrieval, and clustering. Because models can perform effective disambiguation even when their late-layer embeddings appear geometrically similar, relying solely on these embeddings may cause systems to understate the model's actual semantic capabilities. The findings suggest that future methods for evaluating model representations should look beyond static geometry and incorporate the model's predictive behavior.