D y V o: Dynamic Vocabularies for Learned Sparse Retrieval with Entities (2024)

Abstract: Learned Sparse Retrieval (LSR) models use vocabularies from pre-trained transformers, which often split entities into nonsensical fragments. Splitting entities diminishes retrieval accuracy and limits the model’s ability to incorporate up-to-date world knowledge not included in the training data. In this work, we enhance the LSR vocabulary with Wikipedia concepts and entities, enabling the model to resolve ambiguities more effectively and stay current with evolving knowledge. Central to our approach is a Dynamic Vocabulary (DyVo) head, which leverages existing entity embeddings and an entity retrieval component that identifies entities relevant to a query or document. We use the DyVo head to generate entity weights, which are then merged with word piece weights to create joint representations for efficient indexing and retrieval using an inverted index. In experiments across three entity-rich document ranking datasets, the resulting DyVo model substantially outperforms several state-of-the-art baselines.

## introduction neural ir methods typically operate in two stages. initially, a set of candidate documents is retrieved using a fast, computationally-efficient first-stage retrieval method that considers sparse or dense vector representations. these candidates are then re-ranked using more computationally-intensive scoring functions, such as those involving crossencoders @xcite @xcite . learned sparse retrieval (lsr) @xcite @xcite is a prominent neural method for first-stage retrieval. lsr encodes queries and documents into sparse, lexicallyaligned representations that can be stored in an inverted index for fast retrieval. lsr offers several advantages over dense retrieval (dr), another 1 code: https://github.com/thongnt99/dyvo is us a member of who? common approach for first-stage retrieval @xcite ). lsr's lexically grounded representations are more transparent, making it easier for users to understand the model and inspect representations for biases @xcite . furthermore, lsr's compatibility with an inverted index enables efficient and exact retrieval @xcite , while also simplifying the transition from existing lexical search infrastructure supporting methods like bm25. lsr not only performs competitively with dr in terms of performance within the same domain, but it also tends to generalize better across different domains and tasks @xcite . however, lsr models lack explicit representations for entities and concepts in their vocabulary. this can pose challenges due to the tokenization process, where words are segmented into subwords or wordpieces. for instance, a word like "bion-tech" might be tokenized into @xcite . such fragmentation can lead to ambiguity, complicating the retrieval process by obscuring the full meaning and context of the original word, which in turn may affect the accuracy and relevance of search results. additionally, the bag of word pieces representation employed by lsr methods strug-gles with homonyms, where different meanings or entities, such as "who" (world health organization) and "us" (united states), could be conflated when represented merely as word pieces in a query like "is the us a member of who?" hence, while lsr provides a framework for efficient first-stage document retrieval, its current design -particularly in handling entities and complex vocabulary -poses significant challenges. we hypothesize that integrating explicit entities into the lsr vocabulary could significantly enhance its performance. this integration is especially pertinent as a large proportion of queries pertain to specific entities or are closely related to them @xcite . previous work indicates that hybrid models combining word and entity representations have improved both sparse retrieval @xcite @xcite and dense retrieval @xcite @xcite . to address the above limitations, we incorporate entities from wikipedia into the vocabulary of an lsr model. the english wikipedia contains entities spanning a diverse range of categories and disciplines, including named entities like people, organizations, and locations, as well as general concepts such as eudaimonia, hot dog, and net income. integrating these wikipedia entities into a lsr model significantly enhances its ability to handle complex semantic phrases and entities that are currently fragmented into nonsensical word pieces. by enriching query and document representations with relevant entities, we reduce ambiguity and improve the representational power of lsr. this approach is illustrated in figure 1 . moreover, leveraging wikipedia -a rich and continually updated knowledge base -allows the lsr model to refresh its internal knowledge, aligning it with evolving global information. as of april 2024, the english wikipedia hosts nearly 7 million entities and concepts, which is more than 200 times larger than the word piece vocabulary used in current state-of-the-art lsr methods. to identify relevant entities from among millions of them, we propose adding a dynamic vocabulary (dyvo) head with an entity candidate retrieval component. specifically, we leverage entity retrieval techniques and large language models (llms) to dynamically generate relevant entities. these methods aim to refine the set of highly rel-evant entities, which are then passed to the lsr encoder for scoring. the encoder outputs a small bag of weighted entities, ignoring those that were not retrieved. the entity representation is then concatenated with the word-piece representation, forming a joint representation used for indexing and retrieval processes. our contributions are: one of the early works in this area proposed using neural networks to learn sparse representations that are compatible with an inverted index and demonstrated promising performance @xcite . with the advent of the transformer architecture @xcite , subsequent work has successfully utilized pretrained transformers to enhance the effectiveness and efficiency of lsr models @xcite @xcite @xcite . among these, splade @xcite stands out as a state-of-the-art lsr method. while splade uses a word piece vocabulary, prior work has demonstrated that its vocabulary can be replaced by performing additional masked language modeling (mlm) pretraining and then exhaustively scoring all terms in the new vocabulary @xcite . in this work, we dynamically augment a word piece vocabulary using pre-existing embeddings rather than performing additional pretraining. splade typically employs a shared mlm encoder for both queries and documents, enabling term expansion and weighting on both sides. however, previous work @xcite has shown that removing query expansion by replacing the mlm query encoder with an mlp encoder can simplify training and improve efficiency by reducing the number of query terms involved. while most lsr research has focused on ad-hoc paragraph retrieval tasks, recent efforts have explored extending lsr to other settings, such as conversational search @xcite , long documents @xcite , and text-image search @xcite @xcite . entity-oriented search. early work in entityoriented search primarily utilized entities for query expansion. a significant advancement in this domain was made by @xcite , who introduced a double translation process where a query was first translated into relevant entities, and then the terms associated with these entities were used to expand the query. @xcite further developed this concept with entity query feature expansion, which enhanced document retrieval by enriching the query context with entity features. the field then recognized the more integral role of entities in search applications, transitioning from merely using entities for query expansion to treating them as a latent layer while maintaining the original document and query representations. among these methods, explicit semantic analysis @xcite used "concept vectors" from knowledge repositories like wikipedia to generate vector-based semantic representations. the latent entity space model @xcite utilized entities to assess relevance between documents and queries based on their alignments in the entity-informed dimensions. esdrank @xcite leveraged semi-structured data such as controlled vocabularies and knowledge bases to connect queries and documents, pioneering a novel approach to document representation and ranking based on interrelated entities. this progression in research inspired a shift towards methodologies that treated entities not just as a latent layer but as explicit, integral elements of the retrieval model. for example, the creation of entity-based language models marked a significant development. @xcite explored the impact of explicit entity markup within queries and documents, balancing term-based and entitybased information for document ranking. @xcite developed the semantic enabled language model, which ranks documents based on semantic relatedness to the query. xiong et al.'s line of work @xcite @xcite exemplifies a dual-layered approach that pairs a traditional bag of terms representation with a separate bag of entities representation, enhancing the document retrieval process by incorporating both term and entity-based semantics. for example, explicit semantic ranking used a knowledge graph to create "soft matches" in the entity space, and the word-entity duet model captured multiple interactions between queries and documents using a mixture of term and entity vectors. the entity-duet ranking model (edrm) @xcite represents a pioneering effort in neural entity-based search, merging the word-entity duet framework with the capabilities of neural networks and knowledge graphs (kgs). tran and yates (2022) advanced this area by introducing a method that clusters entities within documents to produce multiple entity "views" or perspectives, enhancing the understanding and interpretation of various facets of a document. recently, chatterjee et al. (2024) proposed to learn query-specific weights for entities within candidate documents to re-rank them. entity ranking. the task of entity ranking involves retrieving and ordering entities from a knowledge graph based on their relevance to a given query. traditionally, this process has utilized term-based representations or descriptions derived from unstructured sources or structured knowledge bases like dbpedia @xcite . ranking was commonly performed using models such as bm25 @xcite . additionally, markov randon fields-based models like the sequential dependence model @xcite and its variants @xcite @xcite addressed the joint distribution of entity terms from semi-structured data. as the availability of large-scale knowledge graphs increased, semantically enriched models were developed. these models leverage aspects such as entity types @xcite @xcite and the relationships between entities @xcite to enhance ranking accuracy. more recently, the focus has shifted towards learning-to-rank (ltr) methods @xcite @xcite , which utilize a variety of features, particularly textual information and neighboring relationships, to re-rank entities. the introduction of graph embedding-based models like geeer @xcite and kewer @xcite has further enriched the field by incorporating wikipedia2vec @xcite embeddings, allowing entities and words to be jointly embedded in the same vector space. the latest advancements in this domain have been driven by transformer-based neural models such as genre @xcite , bert-er++ @xcite , and em-bert @xcite . these models introduce sophisticated techniques including autoregressive entity ranking, blending bert-based entity rankings with additional features, and augmenting bert @xcite with wikipedia2vec embeddings. ## methodology in this section, we describe our approach to producing sparse representations of queries and documents that contain both entities and terms from a word piece vocabulary. to do so, we incorporate entities into the model's vocabulary through the use of a dynamic vocabulary head. ## experimental setup datasets. given our need for entity-rich queries and documents, we evaluate our approach using datasets containing a mix of news documents and complex information needs (i.e., trec robust04, trec core 2018, codec), which have also been commonly used in prior work, e.g. @xcite @xcite @xcite . robust04 @xcite has 528k documents and 250 query topics where documents are news articles. all topics are deeply annotated with 1246 judged documents per topic on average. core 2018 contains 595k news articles or blog posts from the washington post with about 50 topics and 524 relevant judgements per topic. codec @xcite provides 729k web documents crawled from various sources and 42 complex query topics, covering recent themes (e.g., bitcoins, nft) from diverse domains, such as history, economics, politics. furthermore, each topic comes with approximately 147 document judgements and 269 entity annotations. we use all provided topics (description field on trec datasets and query field on codec) for evaluation. to train models, we used synthesized dataset provided by inparsv2 @xcite . because codec is not available on in-parsv2, we generate 10k queries ourselves using mixtral-8x7b-instruct-v0.1 @xcite . knowledge base and entity candidates. we use the kilt @xcite knowledge repository with 5.9 millions entities and only keep entities appearing in wikipedia2vec @xcite , resulting in ~5.3 millions entities. to obtain linked entity candidates for queries, we use the rel (van hulst et al., 2020) entity linker with n-gram ner tagger. for the entity retrieval approach, we experimented with different aproaches, including traditional sparse retrieval (bm25), dense retrieval (laque), and generative retrievers @xcite . for bm25, we index the entity's description and retrieve the top 20 entities per query using pyserini @xcite with the default parameters. with laque, we encode both queries and entity descriptions using the laque (distilbert) dense encoder, and select the top 20 entities that have the highest dot product with the query's dense vector. with generative approaches, we prompt mixtral and gpt4 to generate relevant entities and remove out-of-vocabulary entities. for simplicity, we re-use the linked entities from @xcite on the document side for all experiments. training configuration. starting from a lsr checkpoint without entities trained on msmarco, we further fine-tune them on the three datasets using the synthesized queries, monot5-3b scores for distillation, kl loss @xcite and bm25 negatives. to regularize vector sparsity, we apply a l1 penalty on the output sparse representations, which has previously been shown to be effective @xcite . we experiment with different l1 weights, including [1e-3, 1e-4, 1e-5]. for each setting, we train two lsr versions: lsr-w that produces word piece representations only, and dyvo that produces joint word-entity rep-resentations. on each dataset, we train the models for 100k steps with a batch size of 16, learning rate of 5e-7, and 16-bit precision on a single a100 gpu. entity embeddings are pre-computed and frozen during training; only a projection layer is trained where the word and entity embedding sizes differ. ## experimental results we first consider whether incorporating linked entities in sparse representations increases effectiveness over representations containing only word pieces, finding that doing so yields consistent improvements on our entity-rich benchmarks. we then consider the impact of the entity selection component and the entity embeddings used, finding that performing entity retrieval rather than entity linking can further improve performance and that dyvo performs well with a range of entity embedding techniques. rq1: does incorporating linked entities improve the effectiveness of lsr? in this rq, we seek to evaluate the effectiveness of lsr with linked entities. we train three different lsr versions with different sparse regularization weights (1e-3, 1e-4, 1e-5). for each lsr version, we trained two models (lsr-w and dyvo) with and without entities, respectively, using exactly the same training configuration. although we are mainly interested in the comparison between dyvo and lsr-w, other baselines (e.g., bm25, bm25+rm3, and zero-shot single-vector dense retrieval methods) are provided in our first observation is that our model with linked entities (dyvo) outperforms the model without entities (lsr-w) consistently on all metrics (ndcg@10, ndcg@20, r@1000) across three different datasets and three different sparsity constraints. the difference between the two models is more pronounced when the document representations become more sparse. with the largest regularization weight (@xmath0, the documents are the most sparse and have the fewest terms. in this scenario, enriching the word representation with linked entities typically results in a significant gain, notably with an increase ranging from 1.15 to 3.57 points in ndcg@10 across all datasets. when we relax the sparsity regularization to 1e-4 and 1e-5, we observe an improvement in the performance of lsr-w baseline models. however, we still consistently observe the usefulness of linked entities, albeit to a lesser degree. in the most relaxed setup (@xmath1, we often gain from 1 to 2 ndcg points on all three datasets. the r@1000 improvement is similar, except we only observe a minimal increase on core 2018. compared to other families, both lsr and dyvo demonstrate better performance than unsupervised lexical retrieval methods (bm25, bm25+rm3) and dense retrieval (dr) models, including distilbert-dot-v5 @xcite , gtr-t5-base @xcite , and sentence-t5-base @xcite . despite using models three times larger (t5-base vs. distil-bert), both gtr-t5-base and sentence-t5-base still show lower effectiveness than lsr models. this is due to the generalization difficulties of dense retrieval methods. dyvo also outperforms bm25 + rm3, a traditional query expansion method using pseudorelevance feedback. compared to grf, a llmbased query expansion approach by @xcite , dyvo achieves a significantly higher ndcg@10 score (e.g., 53.40 with dyvo using the rel linker versus 40.50 with grf on codec). it is important to note that the llms used in grf were not fine-tuned, and doing so would present substantial computational challenges. rq2: can lsr be more effective with retrieval-oriented entity candidates? in the previous rq, we explored how incorporating linked entities enhances lsr's representations. however, relying solely on linked entities overlooks other relevant entities crucial for document retrieval. for instance, with the codec query "why are many commentators arguing nfts are the next big investment category?", entities like "cryptocurrency", "bitcoin", and "digital asset" can be valuable despite not being explicitly mentioned. our generative approach utilizing mixtral and gpt4 represents a significant step forward in entity retrieval for document ranking. compared to linked entities provided by rel, our approach showcases notable improvements, enhancing ndcg@10 and ndcg@20 scores by approximately +1.3 to +1.78 points across all datasets, with the exception of ndcg@10 on core 2018. mixtral's effectiveness is further highlighted by its impact on r@1000 scores, with increases observed across the ro-bust04, core 2018, and codec datasets. additionally, when we replace mixtral with gpt4, we see further improvements that result in dyvo achieving the highest performance on nearly every metric and dataset. notably, retrieval using gpt-4 generated entities is competitive with re-trieval using human-annotated entities on codec, underlining the significance of enriching query representations with relevant entities beyond linked ones. we attach examples in ## conclusion lsr has emerged as a competitive method for firststage retrieval. in this work, we observed that relying on only word pieces for lexical grounding can create ambiguity in sparse representations-especially when entities are split into subwords. we explored whether learned sparse representations can include entity dimensions in addition to word piece dimensions. in order to facilitate modeling millions of potential entities, we proposed a dynamic vocabulary (dyvo) head that leverages entity retrieval to identify potential entity candidates and entity embeddings to represent them. we find that while both linked entities and llm-generated entities are effective, llm-generated entities ultimately yield higher lsr effectiveness. the approach is largely robust to the choice of entity embedding. our work sets the stage for other lsr models that go beyond word piece vocabularies.