Alleviating Performance Degradation Caused by Out-of-Distribution Issues in Embedding-Based Retrieval (2025)

Abstract: In Embedding Based Retrieval (EBR), Approximate Nearest Neighbor (ANN) algorithms are widely adopted for efficient large-scale search. However, recent studies reveal a query out-of-distribution (OOD) issue, where query and base embeddings follow mismatched distributions, significantly degrading ANN performance. In this work, we empirically verify the generality of this phenomenon and provide a quantitative analysis. To mitigate the distributional gap, we introduce a distribution regularizer into the encoder training objective, encouraging alignment between query and base embeddings. Extensive experiments across multiple datasets, encoders, and ANN indices show that our method consistently improves retrieval performance.

## introduction embedding-based retrieval (ebr) leverages deep encoders, especially pre-trained language models (plms) @xcite , to convert text into high-dimensional embeddings @xcite , enabling fast similarity search via metrics like inner product or euclidean distance. ebr has been widely used in web search @xcite , recommendation @xcite , qa @xcite , and dialogue systems @xcite . to support large-scale scenarios with billions of items, modern retrieval systems combine ebr with approximate nearest neighbor (ann) algorithms @xcite , achieving sub-linear search time with acceptable accuracy. training objectives in ebr typically aim to enhance retrieval accuracy by separating positive and negative examples in the embedding space @xcite . contrastive learning with bi-encoder models @xcite is widely adopted: it pulls positive pairs closer and pushes negatives apart, leading to more discriminative embeddings. this paradigm has proven effective across many plm-based retrievers @xcite @xcite . the trained encoder supports efficient online search by embedding queries for approximate nearest neighbor (ann) based retrieval over a precomputed candidate index @xcite . while contrastive training improves discriminative power between positives and negatives, it also amplifies the distributional gap between query and base embeddings. this issue, first identified in multi-modal training @xcite and later observed in text-only settings @xcite , causes query embeddings to become out-of-distribution (ood) relative to the base data. such distributional mismatch violates the assumption-underlying most ann algorithms-that query and base embeddings follow the same distribution, leading to degraded retrieval accuracy @xcite . in practice, retrieval encoders are often trained without considering ann-specific constraints, under the assumption that embeddings optimized for knn will generalize to ann. however, this overlooks the impact of ood queries, resulting in substantial yet under-recognized performance drops. recent methods like ood-diskann @xcite and roargraph @xcite improve ann robustness by incorporating query information during index construction. though effective, these methods focus on the index side and leave the root cause-ood query embeddings-largely unaddressed. this work empirically shows that query embeddings in retrieval tasks are often ood relative to base embeddings, degrading ann performance. we employ maximum mean discrepancy (mmd) @xcite to quantify the discrepancy, and derive a training-time regularizer to reduce the discrepancy and enhance retrieval robustness. unlike prior work @xcite ) that improves ann indices to tolerate ood queries, we aim to mitigate the root cause by reducing the ood effect at the embedding level during encoder training. to our knowledge, this is among the first works to explicitly target this overlooked distribution gap in an end-to-end text-to-text retrieval setting, where query ood has critical impact on ann accuracy. contributions. we summarize our main contributions as: ## ood in ann search retrieval task setting. in embedding-based retrieval, query q and base data y are encoded by e q and e y into embeddings @xmath0, with distributions p q and p y . similarity sim(•, •), usually inner product, measures their closeness. the training of the encoders involves triplets {q, y + , y -} with positive y + and negatives y -, optimized by infonce loss @xcite : the out-of-distribution (ood) issue. our definition of ood follows @xcite , where query embeddings q are ood relative to base embeddings y if their distributions differ substantially. this differs from treating individual outliers as ood; here, ood describes the entire query embedding distribution p q lies outside the base embedding distribution p y . to quantify this discrepancy, we use the squared maximum mean discrepancy (mmd) @xcite between the query and base embedding distributions, denoted as: based on the base embedding distribution, we compute single-point mmd scores for each query embedding. as shown in figure 1 , in-distribution (id) embeddings (sampled from the base set) yield lower scores, while ood embeddings exhibit significantly higher values, confirming the presence of ood in text-to-text retrieval consistent with prior ann observations @xcite . challenges of serving ood queries in ann search. ood queries are challenging to serve in ann searches because they break the fundamental assumption that queries and base data share the same distribution, leading to significantly larger search spaces @xcite . this results in inefficient search convergence, requiring more computations, memory access, and extended search paths. to demonstrate this issue, we evaluated two ann algorithms, ivf and hnsw, on the ms-marco dataset @xcite using embeddings from cocondenser @xcite and retromae @xcite . ood queries were taken from msmarco dev set, while the ann indices were built on the base embeddings. we compared the search complexity-measured by candidates accessed-for id and ood queries at equal recall. as shown in figure 2 , ood queries require significantly higher complexity, resulting in longer search times and making it difficult to achieve satisfactory recall efficiently. this suggests ann indexes constructed on base distributions poorly generalize to ood queries. contrastive learning amplifies distribution gap. we empirically observe that contrastive learning with infonce increases the distance between query and base embeddings. during finetuning on msmarco with cocondenser and retromae, we monitor the average ℓ 2 distance between query-positive pairs. since the squared ℓ 2 distance can serve as a proxy measure for mmd with a linear kernel (proof in appendix a), it enables us to track distributional shift more conveniently. as shown in figure 3 , the ℓ 2 distance steadily increases during training, indicating that infonce implicitly enlarges the query-base distribution gap. ## solution to the ood issue section 2 highlights the ood issue in retrieval and its harm to ann search, with the squared ℓ 2 distance showing a growing trend as a proxy for distributional drift. in this section, we propose a method to mitigate ood by explicitly regularizing the drift via instance-level decomposition of mmd. for a single query embedding q ∼ p q and its corresponding document embeddings {y + ∪ y -} ∼ p y , we construct the regularization loss directly from the mmd decomposition: ## experiments in this section, we evaluate our method on two standard retrieval benchmarks, msmarco @xcite and natural questions @xcite , measuring recall versus queries per second (qps) following @xcite . experiments use two strong plm-based encoders, cocondenser @xcite and retro-mae @xcite . ann search is conducted with five algorithms: ivf, ivfpq, hnsw (via faiss), diskann, and roargraph @xcite . the distribution regularization weight λ is set to 0.01 to balance with contrastive loss, and training follows official fine-tuning setup. more implementation details are provided in appendix e. ## conclusion we identify and quantify the query ood issue in embedding-based retrieval and show its adverse impact on ann performance. to address this, we propose a simple ℓ 2 -based regularization that narrows the query-base embedding gap with-out adding inference-time overhead. experiments across datasets, encoders, and ann methods show consistent recall gains, especially in high-qps regimes. as future work, we plan to include direct comparisons with recent embedding models such as gritlm @xcite , and to conduct a more systematic analysis of training and inference efficiency.