On the Representations of Entities in Auto-regressive Large Language Models

TL;DR

This paper probes how named entities are represented inside autoregressive LLMs by introducing entity mention reconstruction and the Entity Lens, an extension of the logit lens to decode multi-token mentions. Using task vectors trained per layer, the authors demonstrate that entities are encoded in distributed representations that can be decoded into complete mentions, with context enabling near-perfect copying. They show that entity representations are influenced by frequency rather than token length, and that averaging or linear cleaning can yield better decodability than last-token embeddings. The work further extends to relation decoding, showing that simple linear transformations can extract relational signals between entities, and provides an actionable visualization tool for inspecting what entities the model associates at different layers. Overall, the findings reveal an entity representation space within LLMs, with implications for interpretability, knowledge retrieval, and controllable generation.

Abstract

Named entities are fundamental building blocks of knowledge in text, grounding factual information and structuring relationships within language. Despite their importance, it remains unclear how Large Language Models (LLMs) internally represent entities. Prior research has primarily examined explicit relationships, but little is known about entity representations themselves. We introduce entity mention reconstruction as a novel framework for studying how LLMs encode and manipulate entities. We investigate whether entity mentions can be generated from internal representations, how multi-token entities are encoded beyond last-token embeddings, and whether these representations capture relational knowledge. Our proposed method, leveraging _task vectors_, allows to consistently generate multi-token mentions from various entity representations derived from the LLMs hidden states. We thus introduce the _Entity Lens_, extending the _logit-lens_ to predict multi-token mentions. Our results bring new evidence that LLMs develop entity-specific mechanisms to represent and manipulate any multi-token entities, including those unseen during training. Our code is avalable at https://github.com/VictorMorand/EntityRepresentations .

Figures (18)

• Figure 1: Representation extraction principle: the representation ${\boldsymbol{z}}^\ell$ (in green) of the last token of a target entity mention - here tower for the named entity Eiffel tower - is extracted at a given layer $\ell$ of the transformer model.
• Figure 2: Our entity mention decoding method, in both uncontextual (left) and contextual (right) senarii.
• Figure 3: Distribution of the number of tokens for named entity mentions in the test set (Pythia tokenizer). 87% of them are tokenized with two or more tokens. The dashed line is the median token count (3).
• Figure 4: Context improves decoding. Better performances are obtained on representations extracted in middle layers. Curves present the rate of exact match (on $y$-axis) after training a task vector on representations extracted at layer $\ell$ ($x$-axis is normalized layer $\ell/N_L$, as model have a different number of layers $N_L$).
• Figure 5: Uncontextual mention generation performance is higher for more frequent entities. Performance is analyzed by entity length and mention frequency in the Pilegao2020pile800gbdatasetdiverse. For each model, we chose the layer with best exact match on the \ref{['subsec:Dataset']}. Empty cells indicate fewer than five samples. See \ref{['sec:perfNtokenFreq']} for full results across models and settings.