Table of Contents
Fetching ...

Influence Guided Sampling for Domain Adaptation of Text Retrievers

Meet Doshi, Vishwajeet Kumar, Yulong Li, Jaydeep Sen

TL;DR

This paper tackles how to sample from diverse training corpora to adapt text retrievers to target domains. It introduces Inf-DDS, a bilevel reinforcement-learning framework that uses online influence scores as rewards to learn a data-sampling policy, while reusing computations via a Reptile-style update to keep compute efficient. Across BEIR, MLDR, and Sentence-Transformers datasets, Inf-DDS delivers superior or competitive improvements in NDCG@10 compared with static, gradient-based, and clustering baselines, and exhibits stable, interpretable sampling trajectories. The approach reduces GPU hours relative to some baselines and demonstrates strong capacity for cross-domain and multilingual adaptation, pointing to practical gains for large-scale dense retrieval systems.

Abstract

General-purpose open-domain dense retrieval systems are usually trained with a large, eclectic mix of corpora and search tasks. How should these diverse corpora and tasks be sampled for training? Conventional approaches sample them uniformly, proportional to their instance population sizes, or depend on human-level expert supervision. It is well known that the training data sampling strategy can greatly impact model performance. However, how to find the optimal strategy has not been adequately studied in the context of embedding models. We propose Inf-DDS, a novel reinforcement learning driven sampling framework that adaptively reweighs training datasets guided by influence-based reward signals and is much more lightweight with respect to GPU consumption. Our technique iteratively refines the sampling policy, prioritizing datasets that maximize model performance on a target development set. We evaluate the efficacy of our sampling strategy on a wide range of text retrieval tasks, demonstrating strong improvements in retrieval performance and better adaptation compared to existing gradient-based sampling methods, while also being 1.5x to 4x cheaper in GPU compute. Our sampling strategy achieves a 5.03 absolute NDCG@10 improvement while training a multilingual bge-m3 model and an absolute NDCG@10 improvement of 0.94 while training all-MiniLM-L6-v2, even when starting from expert-assigned weights on a large pool of training datasets.

Influence Guided Sampling for Domain Adaptation of Text Retrievers

TL;DR

This paper tackles how to sample from diverse training corpora to adapt text retrievers to target domains. It introduces Inf-DDS, a bilevel reinforcement-learning framework that uses online influence scores as rewards to learn a data-sampling policy, while reusing computations via a Reptile-style update to keep compute efficient. Across BEIR, MLDR, and Sentence-Transformers datasets, Inf-DDS delivers superior or competitive improvements in NDCG@10 compared with static, gradient-based, and clustering baselines, and exhibits stable, interpretable sampling trajectories. The approach reduces GPU hours relative to some baselines and demonstrates strong capacity for cross-domain and multilingual adaptation, pointing to practical gains for large-scale dense retrieval systems.

Abstract

General-purpose open-domain dense retrieval systems are usually trained with a large, eclectic mix of corpora and search tasks. How should these diverse corpora and tasks be sampled for training? Conventional approaches sample them uniformly, proportional to their instance population sizes, or depend on human-level expert supervision. It is well known that the training data sampling strategy can greatly impact model performance. However, how to find the optimal strategy has not been adequately studied in the context of embedding models. We propose Inf-DDS, a novel reinforcement learning driven sampling framework that adaptively reweighs training datasets guided by influence-based reward signals and is much more lightweight with respect to GPU consumption. Our technique iteratively refines the sampling policy, prioritizing datasets that maximize model performance on a target development set. We evaluate the efficacy of our sampling strategy on a wide range of text retrieval tasks, demonstrating strong improvements in retrieval performance and better adaptation compared to existing gradient-based sampling methods, while also being 1.5x to 4x cheaper in GPU compute. Our sampling strategy achieves a 5.03 absolute NDCG@10 improvement while training a multilingual bge-m3 model and an absolute NDCG@10 improvement of 0.94 while training all-MiniLM-L6-v2, even when starting from expert-assigned weights on a large pool of training datasets.
Paper Structure (29 sections, 7 equations, 17 figures, 14 tables, 1 algorithm)

This paper contains 29 sections, 7 equations, 17 figures, 14 tables, 1 algorithm.

Figures (17)

  • Figure 1: Trade-off between training time (GPU hours) and average NDCG@10 on the MLDR-13 test set when optimizing bge-m3-dense on the MLDR-13 dev set.
  • Figure 2: Overview of Inf-DDS. The trainable scorer $\psi$ and model parameters $\theta$ are optimized by generating online proxy models to compute influence scores, which serve as rewards for updating the scorer $\psi$. Proxy model gradients are efficiently reused for a weighted Reptile update on $\theta$.
  • Figure 3: FEVER training set sampling trajectories for different initialization temperatures using MultiDDS, Inf-DDS, and learned weights from baseline methods (optimized on FEVER dev set).
  • Figure 4: (a) Heatmap showing Z-score (row) normalized performance correlations between train and test splits across BEIR datasets. (b) Domain weights learned by Inf-DDS during optimization for each target domain.
  • Figure 5: Sampling probability initialization vs. Average NDCG@10 on the FEVER and BEIR-5 test set while training with BEIR-7 train set. (a) Scorer optimized only on the FEVER dev set. (b) Scorer optimized jointly on BEIR-5 dev set.
  • ...and 12 more figures