TABi: Type-Aware Bi-Encoders for Open-Domain Entity Retrieval

Findings of the Association for Computational Linguistics: ACL 2022, pages 2147 - 2166 May 22-27, 2022 c⃝2022 Association for Computational Linguistics TABi: Type-Aware Bi-Encoders for Open-Domain Entity Retrieval Megan Leszczynski Daniel Y. Fu Mayee F. Chen Christopher Ré Department of Computer Science, Stanford University {mleszczy, danfu, mfchen, chrismre}@cs.stanford.edu Abstract Entity retrieval—retrieving information about entity mentions in a query—is a key step in open-domain tasks, such as question answering or fact checking. However, state-of-the-art entity retrievers struggle to retrieve rare entities for ambiguous mentions due to biases towards popular entities. Incorporating knowledge graph types during training could help over- come popularity biases, but there are several challenges: (1) existing type-based retrieval methods require mention boundaries as input, but open-domain tasks run on unstructured text, (2) type-based methods should not compromise overall performance, and (3) type-based meth- ods should be robust to noisy and missing types. In this work, we introduce TABi, a method to jointly train bi-encoders on knowledge graph types and unstructured text for entity retrieval for open-domain tasks. TABi leverages a type- enforced contrastive loss to encourage entities and queries of similar types to be close in the embedding space. TABi improves retrieval of rare entities on the Ambiguous Entity Retrieval (AmbER) sets, while maintaining strong overall retrieval performance on open-domain tasks in the KILT benchmark compared to state-of-the-art retrievers. TABi is also robust to incomplete type systems, improving rare entity retrieval over baselines with only 5% type coverage of the training dataset. We make our code publicly available.1 1 Introduction Entity retrieval (ER) is the process of finding the most relevant entities in a knowledge base for a natural language query.2 ER is crucial for open-domain NLP tasks, where systems are provided with a query without the information needed to answer the query (Karpukhin et al., 2020). For instance, to answer the query, “What team does 1https://github.com/HazyResearch/tabi 2We use entity retrieval to refer to the page-level document retrieval setting, where entities correspond to Wikipedia pages. George Washington play for?” an open-domain system can use an entity retriever to find infor- mation about George Washington in a knowledge base. Retrieving the correct George Washington in the query above—George Washington the baseball player, rather than George Washington the president—requires the retriever to recognize that keywords “team” and “play” imply George Washington is an athlete. However, recent work has shown that state-of-the-art retrievers exhibit popularity biases and struggle to resolve ambiguous mentions of rare “tail" entities (Chen et al., 2021). The goal of our work is to improve rare entity re- trieval for open-domain NLP tasks. Rare entities are challenging to retrieve when they share a name with more popular entities. For instance, in a sample of Wikipedia, mentions of George Washington refer to the president 93% of the time, so a retriever can do very well by learning a popularity bias and returning thepresidentwheneveritsees“GeorgeWashington.” This strategy performs poorly on rare entities like George Washington the baseball player. To retrieve a rare entity instead of a popular entity for an am- …

Content truncated to 500 words