NER Retriever: Zero-Shot Named Entity Retrieval with Type-Aware Embeddings (2025)

Abstract: We present NER Retriever, a zero-shot retrieval framework for ad-hoc Named Entity Recognition (NER), where a user-defined type description is used to retrieve documents mentioning entities of that type. Instead of relying on fixed schemas or fine-tuned models, our method builds on pretrained language models (LLMs) to embed both entity mentions and type descriptions into a shared semantic space. We show that internal representations—specifically, the value vectors from mid-layer transformer blocks—encode fine-grained type information more effectively than commonly used top-layer embeddings. To refine these representations, we train a lightweight contrastive projection network that aligns type-compatible entities while separating unrelated types. The resulting entity embeddings are compact, type-aware, and well-suited for nearest-neighbor search. Evaluated on three benchmarks, NER Retriever significantly outperforms both lexical (BM25) and dense (sentence-level) retrieval baselines, particularly in low-context settings. Our findings provide empirical support for representation selection within LLMs and demonstrate a practical solution for scalable, schema-free entity retrieval.

## introduction named entity recognition (ner) is a foundational task in natural language processing (nlp), aimed at identifying and classifying entity mentions in text into predefined categories. while traditional ner systems have made substantial progress, they typically rely on extensive labeled data and are constrained to a fixed set of entity types such as person, organization or location. this limits their applicability in real-world settings, where the types of entities of interest can be diverse, domain-specific, and defined in an ad-hoc manner. to address this limitation, fine-grained entity typing @xcite assigns open-ended fine-grained entity types to entities. recently, it was proposed to extend named-entity recognition to named entity retrieval @xcite . in the retrieval task rather than predicting entity types from a closed schema, the goal is to store documents in a way that enables retrieving all text segments that mention entities of a type defined at query time (e.g., "retrieve all texts mentioning a dinosaur or a politician"). this setup reflects practical needs in information retrieval, question answering, and knowledge base construction. however, it poses unique challenges: entity types may be unseen during training, and representations must generalize across fine-grained, open-ended categories. in this work, we propose ner retriever, a zero-shot retrieval framework for ad-hoc named entity retrieval. our method leverages the internal structure of large language models (llms) to obtain type-aware representations of entity mentions. specifically, we extract contextual embeddings from an intermediate transformer layer (e.g., layer 17 of llama 3.1 8b) and distill them into a compact, discriminative embedding space using a contrastive learning objective. this allows the system to generalize to novel entity types without requiring task-specific llm finetuning. figure 1 illustrates a typical use case. during offline indexing, all entity mentions in the corpus (e.g., "microraptor", "lusotitan") are projected into an embedding space where they cluster near their type description (e.g., "dinosaur") and remain distinct from unrelated concepts (e.g., "the wyoming dinosaur museum"). at query time, user-specified type queries are mapped into the same space and used to retrieve relevant documents via nearest-neighbor search. we evaluate ner retriever across three benchmarks: few-nerd @xcite , multi-coner 2 @xcite , and neretrieve @xcite , and find that it outperforms both lexical (bm25) and dense retrieval baselines in most cases. our results demonstrate that the system enables zero-shot named entity retrieval of highly fine-grained entity types across diverse domains, as represented in the evaluation benchmarks. our results highlight the value of mid-layer llm representations and contrastive tuning for zero-shot entity retrieval, and offer insights into the internal structure of type-sensitive knowledge within pretrained language models. ## task formulation and problem definition we follow the task formulation proposed by @xcite , which defines named entity retrieval as retrieving texts that mention entities be-longing to a given entity type. unlike traditional ner, which detects and classifies entity spans in text, this task frames the entity type itself as an open-ended user-provided query (e.g., "mountain", "politician", "dinosaurs") and aims to return all documents or sentences containing mentions of entities of that type. formally, given a query q type representing a specific named entity type, and a corpus @xmath0, d 2 , . . . , d n } of documents, the task is to return a ranked list @xmath1, d j , . . . , d k ] of documents from c. each document in l must contain at least one named entity mention whose type corresponds to the query q type . the task is evaluated under an exhaustive retrieval setting, where systems are expected to return all relevant documents for each query type from a large corpus. the task is designed to support zero-shot retrieval, a setup that reflects real-world scenarios in which users may search for novel or emerging entity categories, and systems must generalize beyond memorized types to retrieve relevant content. ## method overview we adopt an embedding-based approach in which both entity mentions and type queries are embedded into a shared semantic space, enabling similarity-based search over their vector representations. a central design choice is what to embed and how to embed it effectively. while previous approaches have relied on sentence-level representations @xcite or prompted classification @xcite , we argue that directly embedding individual entity mentions offers greater precision and flexibility. this approach allows the system to focus on localized signals, rather than diluting entity-type cues across entire documents. although embedding every entity increases the number of stored vectors, we offset this by using low-dimensional representations tailored for named entity retrieval. to obtain high-quality embeddings, we build on the contextual representations generated by llms. however, unlike most prior work that uses toplayer outputs, we systematically investigate which internal layer and subcomponent best captures entity-type information. we find that the value (v) vectors from a mid-layer (e.g., block 17 in llama 3.1 8b) outperform other representations in distinguishing fine-grained types. the next section provides a detailed analysis of llm layer selection and motivates our choice of representation through a targeted evaluation of type-discriminative power. ## contextual entity representations our approach relies on extracting entity mention embeddings from intermediate layers of a pretrained llm. in this section, we investigate which layer and component yield the most type-sensitive representations, that is, those that best distinguish entities based on their semantic class (e.g., drug vs. disease). while llms encode rich contextual information, not all layers capture type semantics equally well. rather than using final-layer outputs, as is common in sentence embedding models, we empirically search for internal representations that maximize type separability. motivation: embedding choice matters. prior work has shown that sentence embeddings, even those trained for retrieval, often underperform on entity-type discrimination tasks @xcite . this occurs despite llms containing detailed entity knowledge @xcite , suggesting that naive pooling or top-layer represen-tations fail to transfer type-specific information. we hypothesize that hidden representations from earlier layers may better isolate entity type information, and that these representations can be exploited for retrieval if appropriately selected. evaluation setup. we use the few-nerd dataset to evaluate type sensitivity across different llm layers and components. for each of 20 sampled fine-grained entity types, we collect 20 sentences containing an entity mention of that type. we then form balanced pairs: positive pairs consist of two mentions of the same type, while negative pairs consist of mentions of different types. for each mention, we extract its final-token representation and compute the cosine similarity between mention pairs. we then calculate auc scores for a binary classification task (same-type or not). layer and component sweep. we apply this evaluation procedure to several llms, including llama 3.1 8b, t5 11b, mistral 7b, and gemma 2 7b. for each model, we extract hidden representations from all transformer blocks and subcomponents (e.g., self-attention keys, values, output projections). in total, we analyze 416 distinct representation sources from llama 3.1 alone. we repeat the full evaluation five times with different random seeds and report average auc scores. findings. figure 3 shows a heatmap of auc scores across llama 3.1 layers and subcomponents. we observe substantial variation in performance, with mid-to-late layers outperforming both early and final layers. the highest auc is achieved by the value (v) vectors in the self-attention module of block 17, suggesting that this component encodes particularly strong type-level signals. based on these findings, we use the value (v) vectors from block 17 of llama 3.1 8b as the basis for entity embeddings in our retrieval system, described in the next section. 5 type-aware entity embeddings llm-derived contextual embeddings are highdimensional and not inherently optimized for similarity-based named entity retrieval @xcite . to address this, we introduce a lightweight, supervised contrastive projection module that transforms entity representations into a compact, type-aware embedding space suitable for nearest-neighbor search (type-aware entity embeddings block in figure 2 ). we use a two-layer multilayer perceptron (mlp) trained with a triplet contrastive loss @xcite . each training triplet includes three components: (1) an anchor, which is a natural language type description (e.g., politician); (2) a set of positives, which are entity mentions of the same type (e.g., angela merkel, george w. bush); and (3) a set of negatives, which are mentions from different types (e.g., danube river, mount everest). positive examples are sampled uniformly from entities sharing the same type, while negatives include both randomly selected and hard-negative entities from different types that exhibit lexical similarity to the anchor (e.g., sharing similar surface forms or contextual phrases). this encourages the model to distinguish semantically relevant types from superficially similar distractors. the contrastive objective encourages embeddings of same-type entities to be close together and dissimilar types to be well separated. the resulting representations are low-dimensional, discriminative, and aligned with cosine similarity metrics used during retrieval @xcite . ## results and analysis multiconer 2. this dataset presents the most challenging retrieval scenario due to its short and low-context text segments. ner retriever achieves an r-precision of 0.32-more than three times higher than e5-mistral (0.09) and four times higher than bm25 (0.08). these results indicate that sentence-level embedding models and lexical methods struggle in limited-context settings, while our type-aware embeddings remain effective. few-nerd. ner retriever also leads on few-nerd with an r-precision of 0.34, substantially outperforming both nv-embed v2 (0.04) and e5-mistral (0.08). this confirms our system's ability to generalize across manually annotated, diverse fine-grained entity types. neretrieve test. on the large and diverse neretrieve test set, ner retriever achieves results comparable to nv-embed v2 (0.29 vs. 0.28, difference not statistically significant) and bm25 (0.28). this is the only dataset where our method does not clearly outperform all baselines. we attribute this to the dataset's descriptive nature, as its wikipedia-based source often provides explicit cues about entity types. for example, when querying for "french singer," the term singer is fre-quently stated directly in the text. the fact that bm25 also performs competitively supports this interpretation. entity detection. we report additional metrics (precision@50 and precision@200) in statistical significance. we tested significance using the wilcoxon signed-rank test. for each dataset, we compared our method against the strongest alternative (i.e., the highest-scoring competing model) to assess whether any other system significantly outperforms ours. all significance claims are based on a p-value threshold of 0.05. ## ablation studies and analysis we conduct a series of ablation experiments to isolate the contribution of key design choices in our pipeline: (i) selection of the hidden layer used for representation extraction; (ii) method for selecting the token representation of an entity; and (iii) use of a task-specific projection mlp. all ablation experiments were evaluated on few-nerd. ## conclusion we introduced ner retriever, a zero-shot framework for ad-hoc named entity retrieval based on type-aware embeddings derived from llms. by selecting and refining mid-layer representations with contrastive learning, our method enables accurate and efficient retrieval across diverse, unseen entity types. our results highlight the effectiveness of representation selection and suggest new directions for embedding-based ner systems.