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

arXiv: https://arxiv.org/abs/2408.10613 · HTML: https://mlsanity.com/paper/2408.10613/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 20 Aug 2024.

## Built on by (most cited first, 5 of 5)

- [From Ranking to Selection: A Simple but Efficient Dynamic Passage Selector for Retrieval Augmented Generation](https://mlsanity.com/paper/2508.09497/) (2025)
- [A Survey of LLM × DATA](https://mlsanity.com/paper/2505.18458/) (2025)
- [Don't Retrieve, Generate: Prompting LLMs for Synthetic Training Data in Dense Retrieval](https://mlsanity.com/paper/2504.21015/) (2025)
- [LightRetriever: A LLM-based Text Retrieval Architecture with Extremely Faster Query Inference](https://mlsanity.com/paper/2505.12260/) (2025)
- [Distributionally Robust Optimization with Adversarial Data Contamination](https://mlsanity.com/paper/2507.10718/) (2025)

## Abstract (arXiv)

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.