# Retrieval with Learned Similarities

arXiv: https://arxiv.org/abs/2407.15462 · HTML: https://mlsanity.com/paper/2407.15462/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 42 of its 65 references that are in the map.
- First posted on arXiv 22 Jul 2024.

## Builds on (most cited first, 30 of 42)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Billion-scale similarity search with GPUs](https://mlsanity.com/paper/1702.08734/) (2017)
- [SentencePiece: A simple and language independent subword tokenizer and detokenizer for Neural Text Processing](https://mlsanity.com/paper/1808.06226/) (2018)
- [Improving language models by retrieving from trillions of tokens](https://mlsanity.com/paper/2112.04426/) (2021)
- [Self-Attentive Sequential Recommendation](https://mlsanity.com/paper/1808.09781/) (2018)
- [ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT](https://mlsanity.com/paper/2004.12832/) (2020)
- [Neural Collaborative Filtering](https://mlsanity.com/paper/1708.05031/) (2017)
- [Sigmoid-Weighted Linear Units for Neural Network Function Approximation in Reinforcement Learning](https://mlsanity.com/paper/1702.03118/) (2017)
- [Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphs](https://mlsanity.com/paper/1603.09320/) (2016)
- [Image-based Recommendations on Styles and Substitutes](https://mlsanity.com/paper/1506.04757/) (2015)
- [Session-based Recommendations with Recurrent Neural Networks](https://mlsanity.com/paper/1511.06939/) (2015)
- [ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction](https://mlsanity.com/paper/2112.01488/) (2021)
- [Sentence-T5: Scalable Sentence Encoders from Pre-trained Text-to-Text Models](https://mlsanity.com/paper/2108.08877/) (2021)
- [Large Dual Encoders Are Generalizable Retrievers](https://mlsanity.com/paper/2112.07899/) (2021)
- [Autoregressive Entity Retrieval](https://mlsanity.com/paper/2010.00904/) (2020)
- [AutoInt: Automatic Feature Interaction Learning via Self-Attentive Neural Networks](https://mlsanity.com/paper/1810.11921/) (2018)
- [Transformer Memory as a Differentiable Search Index](https://mlsanity.com/paper/2202.06991/) (2022)
- [Actions Speak Louder than Words: Trillion-Parameter Sequential Transducers for Generative Recommendations](https://mlsanity.com/paper/2402.17152/) (2024)
- [Conditional Computation in Neural Networks for faster models](https://mlsanity.com/paper/1511.06297/) (2015)
- [Breaking the Softmax Bottleneck: A High-Rank RNN Language Model](https://mlsanity.com/paper/1711.03953/) (2017)
- [Neural Collaborative Filtering vs. Matrix Factorization Revisited](https://mlsanity.com/paper/2005.09683/) (2020)
- [Multi-Interest Network with Dynamic Routing for Recommendation at Tmall](https://mlsanity.com/paper/1904.08030/) (2019)
- [Autoregressive Search Engines: Generating Substrings as Document Identifiers](https://mlsanity.com/paper/2204.10628/) (2022)
- [A Neural Corpus Indexer for Document Retrieval](https://mlsanity.com/paper/2206.02743/) (2022)
- [Asymmetric LSH (ALSH) for Sublinear Time Maximum Inner Product Search (MIPS)](https://mlsanity.com/paper/1405.5869/) (2014)
- [Learning Tree-based Deep Model for Recommender Systems](https://mlsanity.com/paper/1801.02294/) (2018)

## Built on by (most cited first, 3 of 3)

- [Request-Only Optimization for Recommendation Systems](https://mlsanity.com/paper/2508.05640/) (2025)
- [MINT: Multi-Vector Search Index Tuning](https://mlsanity.com/paper/2504.20018/) (2025)
- [OneShot: Index-in-Ranking with Neural Scoring for Large-Scale Retrieval](https://mlsanity.com/paper/2607.27475/) (2026)

## Abstract (arXiv)

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.