Task-level Distributionally Robust Optimization for Large Language Model-based Dense Retrieval

TL;DR

The paper addresses how heterogeneous training data distributions affect large language model–based dense retrieval (LLM-DR) fine-tuning. It introduces task-level Distributionally Robust Optimization (tDRO), which end-to-end reweights datasets via a proxy model using an InfoNCE-based loss and a relative loss measurement, then transfers learned domain weights to the LLM-DR fine-tuning phase. The approach separates the DRO optimization from core fine-tuning to accommodate different batching strategies and loss scales, achieving robust domain generalization and substantial data-efficiency gains across multilingual, cross-lingual, and monolingual benchmarks (MIRACL, MKQA, BeIR). Experiments show steady retrieval improvements and up to ~30% reduction in used data, signaling practical benefits for scaling LLM-DR under diverse data regimes. These results highlight a principled path to balancing heterogeneous data contributions to improve universal retrieval performance with large models.

Abstract

Large Language Model-based Dense Retrieval (LLM-DR) optimizes over numerous heterogeneous fine-tuning collections from different domains. However, the discussion about its training data distribution is still minimal. Previous studies rely on empirically assigned dataset choices or sampling ratios, which inevitably lead to sub-optimal retrieval performances. In this paper, we propose a new task-level Distributionally Robust Optimization (tDRO) algorithm for LLM-DR fine-tuning, targeted at improving the universal domain generalization ability by end-to-end reweighting the data distribution of each task. The tDRO parameterizes the domain weights and updates them with scaled domain gradients. The optimized weights are then transferred to the LLM-DR fine-tuning to train more robust retrievers. Experiments show optimal improvements in large-scale retrieval benchmarks and reduce up to 30% dataset usage after applying our optimization algorithm with a series of different-sized LLM-DR models.

Figures (3)

• Figure 1: Task-level Distributionally Robust Optimization for Large Language Model-based Dense Retrieval.
• Figure 2: Different batch sampling strategies and negative types for A) LLM-DR Fine-tuning and B) Distributionally Robust Optimization.
• Figure 3: Weights comparison among baseline, tDRO, and other loss measurement designs.