pEBR: A Probabilistic Approach to Embedding Based Retrieval

TL;DR

This paper introduces pEBR, a probabilistic framework for embedding-based retrieval that dynamically adjusts similarity thresholds to improve recall and precision across different query types.

Contribution

pEBR is the first to model item distribution conditioned on queries, enabling dynamic thresholding for better retrieval performance.

Findings

Significantly improves retrieval precision and recall.

Effectively captures differences between head and tail queries.

Outperforms existing fixed-threshold retrieval methods.

Abstract

Embedding-based retrieval aims to learn a shared semantic representation space for both queries and items, enabling efficient and effective item retrieval through approximate nearest neighbor (ANN) algorithms. In current industrial practice, retrieval systems typically retrieve a fixed number of items for each query. However, this fixed-size retrieval often results in insufficient recall for head queries and low precision for tail queries. This limitation largely stems from the dominance of frequentist approaches in loss function design, which fail to address this challenge in industry. In this paper, we propose a novel \textbf{p}robabilistic \textbf{E}mbedding-\textbf{B}ased \textbf{R}etrieval (\textbf{pEBR}) framework. Our method models the item distribution conditioned on each query, enabling the use of a dynamic cosine similarity threshold derived from the cumulative distribution function (CDF) of the probabilistic model. Experimental results demonstrate that pEBR significantly improves both retrieval precision and recall. Furthermore, ablation studies reveal that the probabilistic formulation effectively captures the inherent differences between head-to-tail queries.