UnIte: Uncertainty-based Iterative Document Sampling for Domain Adaptation in Information Retrieval

Abstract: Unsupervised domain adaptation generalizes neural retrievers to an unseen domain by generating pseudo queries on target domain documents. The quality and efficiency of this adaptation critically depend on which documents are selected for pseudo query generation. The existing document sampling method focuses on diversity but fails to capture model uncertainty. In contrast, we propose Uncertainty-based Iterative Document Sampling (UnIte) addressing these limitations by (1) filtering documents with high aleatoric uncertainty and (2) prioritizing those with high epistemic uncertainty, maximizing the learning utility of the current model. We conducted extensive experiments on a large corpus of BEIR with small and large models, showing significant gains of +2.45 and +3.49 nDCG@10 with a smaller training sample size, 4k on average.

• The paper presents OASiS, a novel framework that leverages both aleatoric and epistemic uncertainty to guide document sampling for domain adaptation in information retrieval.
• It employs BM25-based lexical filtering to remove noisy samples and an iterative epistemic assessment to target knowledge gaps, thereby enhancing nDCG@10 performance.
• Empirical evaluations across BEIR benchmarks demonstrate significant improvements in retrieval accuracy and sample efficiency over previous diversity-based methods.

UnIte: Uncertainty-based Iterative Document Sampling for Domain Adaptation in Information Retrieval

Motivation and Analysis of Prior Sampling Methods

Unsupervised Domain Adaptation (UDA) is critical for generalizing neural retrievers to unseen domains. The prevalent approach leverages pseudo-query generation, coupling unlabeled documents with synthetic queries for retriever adaptation. The principal bottleneck arises in document selection: generating queries for large corpora incurs prohibitive computational costs, necessitating robust sampling policies. Previous methods, including diversity-based sampling with DUQGen, overlook model uncertainty, leading to suboptimal coverage and learning signals. Specifically, DUQGen selects documents in low-density (outlier, noisy) regions and high-confidence regions where the retriever already exhibits strong predictive certainty; both prove uninformative for adaptation and, as demonstrated empirically, substantially degrade nDCG@10 performance.

Figure 1: The scatter plot shows that DUQGen tends to select samples from low-density and high-confidence regions, which are suboptimal for domain adaptation.

The theoretical underpinning draws from the taxonomy of uncertainty in machine learning: aleatoric uncertainty (AU) is inherent data ambiguity, while epistemic uncertainty (EU) captures model ignorance regarding domain-specific phenomena [hullermeier2021aleatoric]. High-AU samples often manifest as corpus outliers, and low-EU corresponds to regions already internalized by the retriever. The proposed paradigm in UnIte, named OASiS, is to filter high-AU (noisy/outlier) documents and prioritize high-EU (knowledge-gap, domain-misaligned) samples, ensuring both informativeness and sample reliability.

OASiS Framework: Algorithm and Uncertainty Estimation

OASiS introduces a two-stage iterative pipeline. First, AU filtering is performed by computing the lexical distance (via BM25) to the k-th nearest neighbor. Documents with high normalized distance are considered low-density and are excluded. This process operates model-independently, avoiding the conflation of data and model uncertainty.

Figure 2: OASiS pipeline overview, showing AU filtering and iterative EU-based sampling-training within the Biomedical domain.

The subsequent loop employs EU estimation per iteration. Here, document embeddings are projected onto the vocabulary space using the MLM head, contrasting token-level probabilities against domain-specific IDF statistics. Documents whose embedding fails to represent high-IDF (domain-important) tokens manifest as high-EU, indicating retriever knowledge gaps. Crucially, as model parameters evolve, EU scores shift, necessitating iterative reassessment and sample selection. Within each cluster (formed for diversity), sampling is balanced with a resampling penalty that dynamically reallocates budget away from previously over-sampled clusters. Maximal Marginal Relevance (MMR) is employed to jointly maximize normalized EU and diversity, ensuring each batch is both informative and representative.

Early Stopping via Uncertainty Plateau

To prevent overfitting and maximize computational sample efficiency, OASiS monitors domain-averaged EU using Exponential Moving Average (EMA). The iterative loop terminates when EU plateaus, corresponding empirically to the peak in retrieval performance.

Figure 3: Relationship between model performance and average EU across iterations—uncertainty minimum aligns with peak nDCG@10.

This unsupervised criterion achieves robust convergence even in the absence of labels, enabling adaptation with as few as 3–5k samples, contrasting with fixed-budget baselines.

Experimental Results and Ablation Studies

Extensive evaluation across five BEIR benchmarks and multiple retrievers (DPR, coCondenser, COCO-DR, Qwen3-Embedding-4B) demonstrates OASiS’s superiority. Gains in nDCG@10 are consistently observed: +2.45 (DPR), +0.75 (coCondenser), +0.26 (COCO-DR), and +3.49 (Qwen3-Embedding-4B) relative to DUQGen with fewer pseudo-queries. Ablations reveal that removing EU sampling or AU filtering incurs significant declines—up to 9 nDCG@10 points—highlighting their complementary necessity.

Figure 4: Ablation results contrasting the impact of with/without AU and EU modules in OASiS.

Additionally, the resampling penalty is instrumental in ensuring balanced coverage, with its omission resulting in over-sampling dominant clusters and degraded retrieval scores. Alternative EU estimation methods (e.g., entropy, MC-dropout) fail to match the domain alignment provided by OASiS’s target-aware EU, as shown experimentally.

Hyperparameter Sensitivity and Sample Efficiency

AU filtering threshold and the balance between uncertainty and diversity are hyperparameters tuned on a held-out development set, demonstrating cross-domain robustness (optimal $Z_{thr}=1.5$, balance $\lambda=0.5$). Sample efficiency is evidenced by early stopping at typically 3–5k documents, yielding higher nDCG@10 per sampled query relative to diversity-based baselines.

Figure 5: The medians of the lexical kNN distances across various k-values are illustrated, supporting AU threshold choice.

Figure 6: Comparison of smoothed uncertainty graph trends for different EMA smoothing settings in early stopping.

Practical and Theoretical Implications

OASiS provides a principled uncertainty-aware sampling strategy, addressing both noise avoidance and knowledge-gap targeting for domain adaptation. Practically, the technique reduces computational overhead and annotation cost, achieving sample-efficient adaptation for neural retrievers under constrained resource budgets. Theoretically, its approach exemplifies distribution-aware domain adaptation, leveraging explicit uncertainty separation to dynamically update sampling policy. Its effectiveness validates the integration of corpus statistics and model states in guiding adaptive training.

OASiS’s paradigm is directly generalizable to other adaptation scenarios: refining curriculum learning in NLP, active data selection under budget constraints, and domain shift handling in retrieval, re-ranking, and cross-lingual IR. The methodology could extend to architectures beyond single-vector retrievers, pending further harmonization with multi-vector and re-ranking objectives.

Conclusion

OASiS establishes a rigorous framework for uncertainty-based iterative document sampling in UDA for retrieval models. By disentangling aleatoric and epistemic uncertainty and iteratively adapting sample selection to evolving model states, it achieves both robust performance and sample efficiency. Its empirical superiority, scalability, and principled domain-adaptive architecture recommend it as a foundation for future data selection and adaptation research, both in IR and broader AI contexts (2604.25142).