TLDR: This research introduces a novel framework, “entity mention reconstruction,” to investigate how Large Language Models (LLMs) internally represent and manipulate named entities. By using “task vectors” to prompt LLMs to generate entity mentions from their internal states, the study reveals that LLMs possess specialized mechanisms for handling entities, including those unseen during training. It demonstrates that entity frequency in training data, rather than token length, is a primary factor in decoding accuracy, and introduces the “Entity Lens” for visualizing an LLM’s entity focus at different processing layers.
Large Language Models (LLMs) have achieved incredible feats in natural language processing, but how they truly understand and manipulate knowledge remains a complex mystery. A new research paper delves into one of the most fundamental building blocks of knowledge in text: named entities.
Authored by Victor Morand, Josiane Mothe, and Benjamin Piwowarski, the study, titled “On the Representations of Entities in Auto-regressive Large Language Models,” introduces a groundbreaking framework to shed light on how LLMs internally represent entities. While previous research often focused on the relationships between entities, this work uniquely investigates the representations of entities themselves, especially multi-token entities like “Eiffel Tower,” which have largely been overlooked by single-token analysis methods.
Unveiling Entity Mechanisms with Entity Mention Reconstruction
The core of this research is a novel method called “entity mention reconstruction.” This framework allows researchers to study whether LLMs can generate entity mentions directly from their internal representations, how multi-token entities are encoded beyond just their last token, and if these representations capture relational knowledge.
The methodology employs “task vectors,” which are essentially learned instructions that prompt the LLM to reconstruct an entity mention from its hidden states. The researchers explored two main settings. First, there is Uncontextualized Decoding, where the model attempts to generate the entire entity mention using only its internal representation, without any surrounding text. This tests how much information about the entity is stored within that representation. Second, there is Contextualized Decoding, where the model is given the original text context along with the internal representation. This setup allows the model to potentially “copy” the mention from the context, revealing how LLMs leverage context for entity handling.
The study utilized the CoNLL-2003 dataset, a standard benchmark for Named Entity Recognition, to ensure a diverse range of entity types, lengths, and frequencies. They experimented with various LLMs from the PYTHIA and PHI families, ranging from 140 million to 7 billion parameters.
Key Discoveries: Specialized Entity Handling and the Role of Frequency
The findings provide strong evidence that LLMs develop specific mechanisms for representing and manipulating entities. In the uncontextual setting, some models (like PYTHIA-6.9B, PHI-2, and PHI-3) could accurately decode up to 65% of entity mentions. Interestingly, representations from the middle layers of the models showed the best performance.
When context was provided (contextual decoding), the performance significantly improved, reaching near-optimal 93% exact match. This suggests a crucial insight: LLMs don’t always store the entire multi-token entity within the representation of its last token. Instead, when context is available, they often use their attention mechanisms to identify and copy the entity directly from the surrounding text.
Another significant finding was that the frequency of an entity mention in the LLM’s training data (like “the Pile” dataset) was a more critical factor for accurate decoding than the number of tokens in the entity itself. More frequent entities were easier to generate.
A control experiment, where the models were asked to decode random sequences of tokens, yielded surprisingly poor results (only 9% exact match uncontextually, 19% contextually). This further supports the idea that LLMs treat named entities as special, meaningful units, developing specialized circuits for their manipulation, rather than just processing arbitrary token sequences.
Also Read:
- Deepening Knowledge Integration: How Semantic-Condition Tuning Enhances LLMs for Knowledge Graph Completion
- Unlocking Logic: How Small Language Models Learn Deductive Reasoning
Beyond the Last Token: Improving Representations and the Entity Lens
The researchers also explored ways to obtain better entity representations. They found that averaging the representations of all tokens within an entity mention improved decoding performance compared to using only the last token’s representation. Additionally, applying a “cleaning” linear layer to the extracted representations further boosted accuracy, suggesting that these transformations can refine the entity-specific features.
The study also confirmed that entity representations can be semantically manipulated to decode basic relations, such as linking landmarks to their home countries. However, averaging or cleaning representations, while good for mention decoding, sometimes degraded the ability to capture these relational semantics.
A practical application introduced by the paper is the “Entity Lens.” Extending the concept of the “logit lens,” the Entity Lens allows researchers to visualize which multi-token entity the model is “thinking” about at any given layer while processing text. For example, when processing “The City of Lights iconic landmark,” the Entity Lens showed the model associating “City of Lights” with Paris and “landmark” with the Eiffel Tower at deeper layers.
This research significantly advances our understanding of how LLMs represent and manipulate knowledge, particularly named entities. The findings pave the way for more interpretable and controllable language models, offering new tools like the Entity Lens for deeper analysis. For more details, you can read the full paper here.
