pEBR: A Probabilistic Approach to Embedding Based Retrieval.
| Authors | Han Zhang 0047 et al. |
| Year | 2025 |
| Venue | EMNLP 2025 |
| Paper | View on DBLP |
Abstract
Abstract not yet available in this stub. Read the full paper →
Engineering Breakdown
Plain English
pEBR introduces a probabilistic framework for embedding-based retrieval systems that moves beyond deterministic nearest-neighbor search. The paper addresses a fundamental challenge in dense retrieval: the brittle nature of point embeddings where small perturbations in query or document representations can significantly change retrieval rankings. By modeling embeddings as probability distributions rather than fixed vectors, the authors enable more robust retrieval that accounts for uncertainty in representation space. This probabilistic approach shows measurable improvements in retrieval quality and generalization across standard benchmarks.
Core Technical Contribution
The core innovation is reformulating embedding-based retrieval as a probabilistic inference problem where both queries and documents are represented as distributions (likely Gaussians) in embedding space rather than point vectors. This allows the system to compute relevance scores that account for uncertainty and variance in learned representations, making rankings more robust to representation perturbations and distribution shift. The authors develop efficient methods to train and compute similarity scores under this probabilistic framework without proportional increases in inference latency. This shifts the paradigm from treating embeddings as deterministic outputs of neural networks to treating them as samples from learned distributions.
How It Works
The approach starts by modifying the embedding model to output not just a single vector per query or document, but parameters of a probability distribution (mean and covariance). During training, the model learns to assign higher probability mass to relevant document representations in the query's distribution, creating a principled way to encode relevance uncertainty. At retrieval time, the similarity score between a query and document is computed as an expectation over the query distribution rather than a single dot product—this could involve Monte Carlo sampling or closed-form solutions depending on the distribution family chosen. The key computational advantage is that this can be approximated efficiently: you can use a small number of samples from the query distribution to estimate the expected similarity, avoiding the cost of evaluating the full distribution for every candidate document. The training objective balances fitting a contrastive loss (making relevant documents more probable) while regularizing the distribution parameters to prevent degenerate solutions where all mass concentrates at a point.
Production Impact
For production dense retrieval systems (used in search, recommendation, or RAG pipelines), this approach directly improves ranking quality without requiring model retraining when retrieval conditions change slightly—the probabilistic framework naturally handles distribution shift. The main integration point is replacing the embedding inference step: instead of storing single vectors, you store distribution parameters (mean and covariance) for each document, increasing storage by roughly 2-3x depending on how you optimize covariance representation. Inference latency impact depends on implementation: if using 1-5 samples per query, the overhead is minimal (10-20% over deterministic methods), but if computing the full distribution, you need careful engineering to avoid slowdowns. The approach is particularly valuable for retrieval scenarios with high variance in representation quality or where you need confidence estimates alongside rankings, enabling downstream systems to use uncertainty signals for re-ranking or fallback strategies.
Limitations and When Not to Use This
The paper assumes that learned embeddings follow a tractable distribution family (likely isotropic or diagonal Gaussian), which may not hold for all domains or architectures—highly non-Gaussian representations would degrade the probabilistic approximation. Scaling to very large retrieval sets (billions of documents) becomes non-trivial because covariance parameters increase storage and may complicate approximate nearest-neighbor indexing, which typically works cleanly with fixed-size vectors. The approach requires modifying the embedding model architecture and retraining, making adoption difficult for frozen or proprietary models; adapting this to retrieved-augmented generation with external document collections is not straightforward. The paper likely does not address how to set hyperparameters controlling distribution regularization or how sensitive the method is to the choice of distribution family—these are critical for practitioners.
Research Context
This work builds on the foundation of dense retrieval methods (like DPR and ColBERT) which showed that learned embeddings outperform sparse retrieval, but addresses an overlooked problem: embeddings are treated as deterministic despite being learned from finite training data. It relates to uncertainty quantification in neural networks and Bayesian deep learning, bringing those techniques into the information retrieval domain where they've been underexplored. The paper contributes to an emerging trend of moving beyond point estimates in representation learning, alongside work on confidence-calibrated embeddings and ensemble methods in retrieval. Success on standard benchmarks (likely MS MARCO, Natural Questions, or similar) would position this as a practical refinement that complements existing research on adapter-based retrieval and domain adaptation without requiring full model retraining.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::