TAB i: T ype-Aware Bi-Encoders for Open-Domain Entity Retrieval (2022)

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 overcome 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 methods 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.

## introduction entity retrieval (er) is the process of finding the most relevant entities in a knowledge base for a natural language query. foot_1 er is crucial for open-domain nlp tasks, where systems are provided with a query without the information needed to answer the query @xcite . for instance, to answer the query, "what team does george washington play for?" an open-domain system can use an entity retriever to find information 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 @xcite . the goal of our work is to improve rare entity retrieval 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 the president whenever it sees "george washington." 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 ambiguous mention, the retriever needs to learn to leverage context cues to overcome the popularity bias. however, existing state-of-the-art retrievers for open-domain tasks (e.g., @xcite , dpr @xcite ) are only trained on unstructured text, making it challenging for them to learn to associate context cues (e.g. "team" and "play") with groups of entities (e.g., athletes). a promising approach to overcome popularity biases is to incorporate types (e.g., athlete or politician) from a knowledge graph into the retriever. a key advantage of types is that contextual cues learned over popular entities can generalize to rare entities of the same types. however, there are several challenges with using types for open-domain retrieval. first, existing methods that use types assume mention boundaries are provided in the input @xcite @xcite what team does george washington play for? ## preliminaries we review the problem setup, task, and bi-encoders. problem setup let @xmath0, @xmath1, @xmath2, and @xmath3. 4 we assume as input a labeled dataset @xmath4,e i ,y i ,t i )} n @xmath5, where n is the number of examples. entity retrieval task given a query q as input, the entity retrieval task is to return the top-k entity candidates relevant to the query from y. since our primary motivation is open-domain nlp tasks, we focus on the page-level document retrieval setting, where we assume that each document corresponds to an entity (e.g., wikipedia page) and that no mention boundaries are provided as input. ## approach tabi leverages knowledge graph types and unstructured text to train bi-encoders for open-domain entity retrieval. tabi takes as input queries and entity descriptions and uses a type-enforced contrastive loss. at inference time, tabi uses nearest neighbor search to retrieve entities. ## retrieval experiments our experiments find that tabi can improve rare entity retrieval for open-domain nlp tasks while maintaining strong overall retrieval performance. ## embedding quality analysis we evaluate how well tabi captures types through embedding visualization, nearest neighbor type classification, and an entity similarity task. embedding visualization we use t-sne to qualitatively evaluate how well bi-encoders cluster entity embeddings by type. figure 2 shows that tabi forms tighter type clusters than blink for five figer types. 13 types are not captured as well when the type is only present in the input and not the loss. this suggests that our type-based loss term helps encode types in the embedding space. type classification to better understand how well embeddings are clustered by type, we evaluate query and entity embeddings using knn classification with @xmath6. 14 we use strict accuracy, loose micro f1, and loose macro f1 metrics for evaluation @xcite . tabi outperforms blink on knn classification over query embeddings on figer and ontonotes, confirming that our loss encourages nearby query embeddings to share the 13 we choose popular types with low overlap in entities. 14 as a query or entity can have multiple types, we cast type classification as a multi-label classification problem. same type ( ## discussion we discuss limitations of tabi. first, we assume a relatively coarse type system is available. to pull together query embeddings of the same type, the type system needs to be sufficiently coarse-grained and the batch size large enough such that multiple examples in a randomly sampled batch have the same type. second, our method is designed for open-domain tasks, which tend to have short queries and strong type disambiguation signals. however, there are disambiguation signals that may be present in queries, such as the existence of a knowledge graph relation between two entities, that tabi does not optimize for learning. to address this, we are interested in incorporating other forms of structured data, including different modalities, into our model as future work. entity typing a task closely related to our work is entity typing, or predicting the set of types for a mention (e.g., @xcite @xcite . a key difference is that entity typing methods often learn explicit type embeddings to perform type classification, whereas tabi only learns query and entity embeddings. entity typing methods could be used to add type labels to the training data as an alternative to tabi's approach that uses a direct knowledge graph type mapping. ## conclusion we introduce a method to train bi-encoders on unstructured text and knowledge graph types through a type-enforced contrastive loss. our loss can improve retrieval of rare entities for ambiguous mentions, while maintaining strong overall performance on open-domain nlp tasks. we hope our work inspires future work on integrating structured data into pretrained models.