Retrieval with Learned Similarities

TL;DR

This paper introduces Mixture-of-Logits (MoL), a universal similarity function approximator, enabling efficient retrieval with learned similarities, achieving state-of-the-art results and significantly improved latency in diverse retrieval tasks.

Contribution

The paper proposes MoL as a flexible similarity function approximator, along with techniques for efficient approximate top-k retrieval with tight error bounds, advancing retrieval methods with learned similarities.

Findings

MoL achieves superior performance across various retrieval scenarios.

Approximate top-k algorithms with MoL outperform baselines by up to 66x in latency.

Mutual information-based load balancing loss improves retrieval effectiveness.

Abstract

Retrieval plays a fundamental role in recommendation systems, search, and natural language processing (NLP) by efficiently finding relevant items from a large corpus given a query. Dot products have been widely used as the similarity function in such tasks, enabled by Maximum Inner Product Search (MIPS) algorithms for efficient retrieval. However, state-of-the-art retrieval algorithms have migrated to learned similarities. These advanced approaches encompass multiple query embeddings, complex neural networks, direct item ID decoding via beam search, and hybrid solutions. Unfortunately, we lack efficient solutions for retrieval in these state-of-the-art setups. Our work addresses this gap by investigating efficient retrieval techniques with expressive learned similarity functions. We establish Mixture-of-Logits (MoL) as a universal approximator of similarity functions, demonstrate that MoL's expressiveness can be realized empirically to achieve superior performance on diverse retrieval scenarios, and propose techniques to retrieve the approximate top-k results using MoL with tight error bounds. Through extensive experimentation, we show that MoL, enhanced by our proposed mutual information-based load balancing loss, sets new state-of-the-art results across heterogeneous scenarios, including sequential retrieval models in recommendation systems and finetuning language models for question answering; and our approximate top-$k$ algorithms outperform baselines by up to 66x in latency while achieving >.99 recall rate compared to exact algorithms.