Dynamic Injection of Entity Knowledge into Dense Retrievers (2025)

Abstract: Dense retrievers often struggle with queries involving less-frequent entities due to their limited entity knowledge. We propose the Knowledgeable Passage Retriever (KPR), a BERT-based retriever enhanced with a context-entity attention layer and dynamically updatable entity embeddings. This design enables KPR to incorporate external entity knowledge without retraining. Experiments on three datasets demonstrate that KPR consistently improves retrieval accuracy, with particularly large gains on the EntityQuestions dataset. When built on the off-the-shelf bge-base retriever, KPR achieves state-of-the-art performance among similarly sized models on two datasets. Models and code are released at https://github.com/knowledgeable-embedding/knowledgeable-embedding.

## introduction language models (lms) struggle to capture lessfrequent or up-to-date entity knowledge @xcite , often resulting in hallucinations @xcite . retrievalaugmented generation (rag), which enhances lms by leveraging external knowledge retrieval, is a promising approach to mitigate this issue. dense retrievers are commonly employed for this purpose; however, because they also rely on lms, they likewise struggle with queries involving less-frequent entities @xcite and often fail to retrieve such knowledge effectively. in this paper, we address this problem by proposing a simple extension to dense retrievers that dynamically injects entity knowledge into their embeddings. specifically, we introduce the knowledgeable passage retriever (kpr), a bertbased dense retriever @xcite ) that integrates entity embeddings adaptable at inference time (see figure 1 ). kpr is intentionally designed with a simple architecture and is trained to attend to entity knowledge based on the context of the input dictionary-based entity linker iphone was developed and released by @xcite ... [mask] was established in 1976 in california... who founded apple? [cls] who founded apple ? [sep] input query context-entity attention layer no-op apple (food) apple inc. bert (dpr) query embedding bert (inference) entity embedder query keys & values ⠇ ⠇ wikipedia passages referring an entity kpr computing relevance via dot product with passage embeddings attending relevant entities based on input context figure 1: architecture of kpr, a dense retriever with a context-entity attention layer that incorporates entity knowledge, given a query as input. entity embeddings are obtained via bert inference, and both the embeddings and the entity linker can be updated without retraining. relevance is scored by the dot product of query and passage embeddings, both produced by kpr. to evaluate kpr, we use the entityquestions (eq) dataset @xcite , which includes many queries with less-frequent entities, as well as the natural questions (nq) @xcite and triviaqa (tqa) @xcite datasets. kpr achieves a substantial 12.6% gain in top 20 accuracy on eq over the model without kpr extensions and consistently improves performance on the other two datasets. we further build kpr on top of the off-the-shelf bge-base retriever, achieving state-of-the-art results on eq and tqa compared to other retrievers of similar size. ## method kpr is built upon dpr @xcite , a widely used bert-based dense retriever. we add an attention layer on top of dpr to incorporate entity knowledge from a knowledge base (kb) (see figure 1 ). we use wikipedia as the target kb. dpr. dpr encodes a query or passage into a ddimensional embedding, obtained from the output embedding of bert's @xcite token. relevance between a query and a passage is scored as the dot product of their embeddings. the model is fine-tuned on datasets containing queries paired with positive and negative passages, using a crossentropy loss with the dot product scores as logits. see appendix a for further details. input entity representation. kpr uses the following two d-dimensional input embeddings: kpr takes the input tokens and a set of entities @xmath0, . . . , e n detected from the text using an entity linker. the input representation of each entity is computed by summing its entity embedding and entity position embedding based on its position in the input token sequence. if an entity spans multiple tokens, its position embedding is computed as the average of the embeddings of the corresponding positions @xcite . to prevent the entity sequence from being empty, a d-dimensional no-op embedding is appended to the sequence of input entity representations. context-entity attention layer. kpr adopts a single-head key-query-value attention mechanism @xcite . the query matrix q is computed based on the output embedding of bert's [cls] token (denoted by h [cls] ∈ r 1×d ). the key and value matrices, k and v, are computed based on the input entity representations denoted by @xmath1: entity embedder. we compute entity embeddings based on bert inference over wikipedia passages that refer to the corresponding entity @xcite . specifically, for each passage containing an anchor link to the entity, we replace the entity name with bert's [mask] token, extract the output embedding of the [mask] token, and average the resulting embeddings across passages. the final embeddings are then normalized to match the average norm of bert's input token embeddings. entity linker. we use dictionary-based string matching to detect entities mentioned in the text. the dictionary comprises anchor names and their possible referent entities, obtained from wikipedia hyperlinks. for ambiguous names (e.g., "apple," which can refer to apple inc. and apple (food)), we do not disambiguate to a single entity, but instead input all possible candidates to maintain recall. training. the parameters in our attention layer, entity position embeddings, and no-op embedding are initialized randomly. entity embeddings are kept frozen during training, enabling them to be added or updated dynamically after training. the dropout probability in the attention layer is set to match that of the other layers. we follow dpr @xcite for the remaining training settings including the loss function. note that entities can be dynamically updated in kpr without retraining, since kpr relies on entity embeddings computed via bert inference, which are kept frozen during training, and on a dictionarybased entity linker, whose entries can be modified. computational overhead. kpr adds entity embeddings to bert, requiring additional storage. however, since the embeddings are used only as input features, they can be implemented as a sparse lookup table and stored not only in gpu memory but also in cpu memory or even on disk, with minimal impact on speed @xcite . furthermore, kpr introduces only a small overhead in total flops, as discussed in appendix b. datasets and metrics. we train our model on the dataset proposed by @xcite , which is based on nq, tqa, webquestions @xcite , and curatedtrec @xcite . we evaluate the models on the eq, nq, and tqa datasets. following @xcite , we report top 20 retrieval accuracy. top 100 accuracies are also provided in appendix d. we use the 21m wikipedia passages released by @xcite as the target passages. baselines. we use vanilla dpr and conventional bm25, which performs robustly on queries with less frequent entities @xcite . base models. we use base-sized bert and retromae @xcite , a state-of-the-art bert-based model pretrained for retrieval. for dpr, we also adopt pelt @xcite , a bert-based model that injects entity knowledge by inserting each entity embedding, placed between the token embeddings of parentheses, immediately after the corresponding entity token. while pelt introduces entity knowledge through input augmentation, kpr incorporates it via a dedicated attention layer. unlike kpr, pelt requires each input entity name to be disambiguated to a single entity, which incurs additional computational complexity. entity embeddings. we use the same entity embeddings, derived from base-sized bert inference ( §2), for both pelt and all kpr models. entity linker. by default, we use our dictionarybased linker ( §2) and additionally employ the stateof-the-art refined linker @xcite . we select refined over other off-the-shelf systems such as blink @xcite , @xcite , and luke @xcite due to its superior performance. we refer to kpr based on bert as kpr bert and name the other models accordingly. further details are provided in appendix c. ## related work the performance of dense retrievers @xcite remains insufficient for queries involving less-frequent entities @xcite , due to the limited knowledge of such entities in lms @xcite . several studies have explored incorporating entity knowledge to enhance lms @xcite @xcite , among which we adopt pelt @xcite as a baseline. entity knowledge has also been leveraged to improve retrieval tasks @xcite @xcite @xcite . however, to our knowledge, only tran and yates (2022) is directly applicable to dense retrieval, integrating wikipedia2vec embeddings into a bertbased retriever. because its entity representations are computed independently of bert, the model cannot learn to select relevant entities based on context. as a result, it fails to improve performance compared to the vanilla model when passages are represented by a single embedding. moreover, its effectiveness on queries involving less-frequent entities remains unexamined, and updating entity embeddings requires retraining wikipedia2vec. in contrast, kpr integrates an attention layer into bert, enabling it to select contextually relevant entities more effectively. this leads to improved performance across multiple retrieval benchmarks. ## conclusions we proposed kpr, a dense retriever with a contextentity attention layer and dynamically updatable entity knowledge. kpr consistently improves performance across benchmarks, particularly on queries involving less-frequent entities. when built on bge-base, it achieves state-of-the-art results on two benchmarks among similarly sized models. future work includes applying kpr to decoder-based retrievers and extending it to kbs beyond wikipedia.