Research map / Large language models / Instruction and fine-tuning

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

AAAI 2025 · Large language models › Instruction and fine-tuning · approach: Efficient LLM fine-tuning systems · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes task-level Distributionally Robust Optimization (tDRO) that learns per-dataset domain weights with a small proxy model and relative loss measurements, then transfers them to LLM-based dense retriever fine-tuning.

Problem

Choosing dataset mixtures and sampling ratios for LLM-based dense retrieval fine-tuning relies on intuition, and standard DRO methods like GroupDRO are incompatible with task-homogeneous batching and with heterogeneous loss scales across datasets.

Approach

Separate DRO weight learning from LLM-DR fine-tuning: a small proxy model (Qwen1.5-0.5B) runs exponentiated-gradient DRO over datasets using loss divided by a reference model's loss (relative loss measurement), and the learned weights are then used to select datasets or reweight sampling for LLM-DR models of various sizes.

Key results

• Llama3-8B with tDRO top-70% dataset selection reaches 66.4 nDCG@10 on MIRACL vs 64.1 uniform baseline (+2.3).
• tDRO sample ratio reweighting with Qwen-1.8B improves MKQA Acc@10 by +5.5 over uniform sampling.
• tDRO top-70% selection uses roughly 70% of the datasets while improving retrieval on MIRACL, MKQA, and BEIR across all six base models.

Measured effects reported in the paper

Surprising findings

• Using the raw proxy loss or loss difference instead of relative loss measurement hurts performance, with weights over-preferring the Yahoo Answers dataset due to its large loss scale.
• Loss reweighting alone improves BEIR but gives no gains on MIRACL and MKQA, suggesting data-sampling changes matter more than loss scaling.

Limitations

• Gains on monolingual BEIR are small or absent for the largest models (Qwen-7B, Mistral-7B, Llama3-8B).
• Requires a trained reference model and a proxy run before the final LLM-DR fine-tuning, adding pipeline cost.

Experimental setup

Datasets: sentence-transformers training data (25 datasets incl. MS-MARCO, MIRACL, Mr.TyDi, DuReader, T2-Ranking, Yahoo Answers)

Benchmarks: MIRACL MKQA BEIR

Metrics: nDCG@10 Recall@100 Accuracy@10 Accuracy@100

Baselines: Uniform sampling LLM-DR (Qwen-0.5B to Llama3-8B) GroupDRO-style worst-case loss BM25 mContriever mE5-large-instruct E5-Mistral

Builds on: GroupDRO InfoNCE RepLLaMA DoReMi

Closest papers by meaning

• Scaling Retrieval-Based Language Models with a Trillion-Token Datastore (NeurIPS 2024 · similarity 0.86)
• LightRetriever: A LLM-based Text Retrieval Architecture with Extremely Faster Query Inference (ICLR 2026 · similarity 0.86)
• Attention in Large Language Models Yields Efficient Zero-Shot Re-Rankers (ICLR 2025 · similarity 0.86)
• RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval (NeurIPS 2025 · similarity 0.86)
• Search for Efficient Large Language Models (NeurIPS 2024 · similarity 0.86)
• Data Mixing Optimization for Supervised Fine-Tuning of Large Language Models (ICML 2025 · similarity 0.86)