Relevance-Based Embeddings: Lightweight Candidate Retrieval via Heavy-Ranker Calls
Abstract
In many machine learning applications, the most relevant items for a query should be efficiently retrieved. The relevance function is usually an expensive similarity model, making the exhaustive search infeasible. A typical solution is to train another model that separately embeds queries and items to a vector space, where similarity is defined via the dot product or cosine similarity. This allows one to search the relevant items through fast approximate nearest neighbor search at the cost of some reduction in quality. To compensate for this reduction, the found items (candidates) are re-ranked by the expensive ranking model. In this paper, we investigate an alternative approach to candidate selection that utilizes the scores of the expensive model to improve the representations of queries and items. The idea is to describe each query (item) by its relevance to a set of support items (queries) and use these new representations to obtain query (item) embeddings. We theoretically prove that such embeddings are powerful enough to approximate any complex similarity model (under mild conditions). We also investigate the choice of support items, which is a crucial ingredient of the proposed approach. The experiments on diverse academic and production datasets illustrate the power of our method.
Lay Summary
Many search and recommendation systems need to quickly find the most suitable items for a user’s query, such as products, documents, or videos. The best models are often very accurate but too slow and expensive to compare a query with every possible item. A common workaround is to use a faster, smaller or simpler model to first retrieve a small set of documents (candidates), then apply the expensive model only to rearrange them. In this work, we propose an efficient way to build the fast candidate-selection model using information from the expensive model. Instead of training the fast model independently, we use the scores produced by the accurate but costly model to construct better representations of queries and items. Each query or item is described by how relevant it is to a carefully chosen set of “support” examples, and these descriptions are then transformed into efficient representations for search. We show theoretically that this approach can approximate expensive models under reasonable assumptions. We also study how to choose the support examples, which is important for performance. Experiments on both academic and real production datasets show that our method improves candidate retrieval quality while keeping search efficient.