Technical Summary: Influence Guided Sampling for Domain Adaptation of Text Retrievers

Problem Statement

General-purpose open-domain dense retrieval systems are typically trained on large, eclectic mixes of corpora and search tasks. A critical challenge in developing robust, generalizable embedding models is determining how to sample these diverse datasets for training. Conventional strategies—such as uniform sampling, proportional sampling based on instance population, or static expert-defined weights—are often suboptimal. They fail to dynamically adapt to the specific informativeness of data subsets required to maximize performance on a target domain. While dynamic, gradient-based approaches (e.g., DDS, DoGE, DoReMi) have been proposed to learn sampling weights online, they face significant limitations: (1) they suffer from instability and high variance due to stochastic gradient estimates used for reward computation, and (2) they incur substantial computational overhead when requiring proxy models or expensive estimators.

Methodology: Inf-DDS

The authors propose Inf-DDS (Influence-guided Dynamic Data Sampling), a novel reinforcement learning-driven framework designed to adaptively reweight training datasets. The core innovation lies in replacing noisy gradient-based rewards with influence-based reward signals derived from online proxy models.

1. Problem Formulation

The task is framed as a bilevel optimization problem. The goal is to learn a parameterized sampling policy $\psi$ that adjusts the probability distribution $P_D(i; \psi)$ over a pool of training datasets $\{D_i^{train}\}$ to maximize model performance on a set of target development datasets $\{D_j^{dev}\}$.

• Inner Loop: Updates model parameters $\theta$ using the current sampling distribution.
• Outer Loop: Updates the scorer parameters $\psi$ to optimize the sampling distribution based on the impact on the development set.

2. Influence-Based Rewards

Unlike prior methods that rely on first-order approximations or small proxy models to estimate data influence, Inf-DDS employs online proxy models to compute exact influence scores.

• Process: At each time step $t$, the model parameters $\theta_t$ are updated using a specific training dataset $D_i^{train}$ for $l$ gradient steps to produce $\theta_{t+1}^i$.
• Metric: The influence is quantified by the change in performance on the development set: $\Delta M_j^i = M(\theta_{t+1}^i; d_j^{val}) - M(\theta_t; d_j^{val})$.
• Reward: The final influence score $I_i$ is the mean of these performance changes across all development datasets. This score serves as the reward signal to upweight datasets that demonstrably improve downstream performance.

3. Optimization and Efficiency

To address computational scalability, Inf-DDS utilizes a Reptile-style first-order meta-update:

• Gradient Reuse: Instead of maintaining separate optimizer states for each dataset, the method computes a weighted average of the updated parameters $\bar{\theta}_{t+1} = \sum p_i \theta_{t+1}^i$ (where $p_i$ are probabilities derived from influence scores) and performs a single update step. This allows the reuse of a single set of intermediate gradients and optimizer states, significantly reducing memory overhead.
• Subsampling: When the number of domains $M$ is large, the scorer $\psi$ is updated on a random subsample $S$ of datasets to reduce per-iteration compute, while maintaining an unbiased estimator for the conditional objective.

Key Contributions

1. Novel Algorithm: The proposal of Inf-DDS, an influence-guided reinforcement learning approach that learns to adjust sampling probabilities across diverse training datasets without relying on noisy gradient estimates for rewards.
2. Computational Efficiency: The method is designed to be lightweight regarding GPU consumption. By reusing computations from the parameterized sampling distribution updates to also update model parameters, it avoids the heavy overhead associated with proxy models in other dynamic sampling methods.
3. Empirical Validation: The authors validate Inf-DDS across multiple benchmarks (BEIR, MLDR, and the Sentence-Transformers embedding dataset), demonstrating that it produces more stable sampling trajectories and achieves superior or comparable performance to gradient-based methods (MultiDDS, DoGE) and proxy-model approaches (DoReMi).

Experimental Results

The efficacy of Inf-DDS was evaluated on three distinct scenarios:

• BEIR (Domain Adaptation): When optimizing for the FEVER development set, Inf-DDS consistently outperformed static sampling and MultiDDS across various initialization temperatures. It demonstrated stable sampling trajectories, correctly prioritizing informative datasets like MS MARCO and FEVER, whereas gradient-based methods exhibited unstable dynamics.
• MLDR (Multilingual Long Document Retrieval): Using the BGE-M3 corpus, Inf-DDS improved the baseline bge-m3-dense model by 5.03 absolute NDCG@10 points. Notably, it successfully upweighted underrepresented languages while avoiding overfitting on dominant languages (English and Chinese), achieving the highest individual-language performance in 8 out of 13 languages.
• Sentence-Transformers Embedding Dataset: In a challenging setting with 1 billion parallel sentences across 32 domains:
  • Starting from a uniform initialization, Inf-DDS achieved performance nearly matching the off-the-shelf Expert model (all-MiniLM-L6-v2), outperforming the uniform baseline by 1.83 points.
  • Starting from Expert-assigned weights, Inf-DDS provided an additional 0.94 absolute NDCG@10 improvement.
  • In contrast, gradient-based methods (MultiDDS, DoGE) failed to improve upon the baselines in these settings.
  • Inf-DDS achieved these gains with 1.5× to 4× less GPU compute compared to methods like DoReMi.

Significance and Claims

The paper claims that Inf-DDS offers a more robust and computationally tractable solution for domain adaptation in text retrieval. By shifting from gradient-derived rewards to influence-based rewards computed via online proxy models, the method stabilizes the learning of sampling weights and provides a more direct estimate of which data subsets lead to the best attainable minima. The authors emphasize that while the algorithm does not always converge to a global optimum, it reliably delivers substantial improvements from reasonable initializations. The work highlights that effective retriever improvement requires moving beyond simple domain similarity or static heuristics, necessitating adaptive strategies that directly optimize for downstream performance metrics.