论文领读｜tDRO:面向大模型稠密检索的任务级分布鲁棒优化

研究背景

随着基础模型的持续进步和数据量的显著增加，以大语言模型(LLM)为基座的稠密检索(Dense Retrieval, DR)模型已展现出卓越的检索性能。大模型稠密检索(LLM-DR)基于大量异质的微调数据集进行训练，这些异质数据集在领域、语言、对称性等方面展现出巨大的差异性。然而，如何确定一个理想的数据分布或配比，以最大化这些数据集联合训练的效果，成为提升大模型稠密检索(LLM-DR)性能的核心挑战。

目前，关于大模型稠密检索(LLM-DR)的数据分布优化(Data Distributional Optimization) 的研究仍处于空白状态。在现有实践中，决定使用哪些数据集、每个数据集使用的比例以及如何用更少的数据实现更优性能，主要依赖于研究人员的经验性判断和反复的实验。但这种方法不仅试错成本高，而且不能穷尽所有可能的数据配比，导致模型鲁棒性不足，往往只能达到次优状态，从而影响大模型稠密检索(LLM-DR)的性能表现。

为了解决上述异质数据集联合训练的分布优化问题，我们提出了一种面向大模型稠密检索(LLM-DR)的任务级分布鲁棒优化(Task-level Distributionally Robust Optimization, tDRO)算法，该算法可以端到端地学习鲁棒的数据分布。在大规模的开源文本检索训练集的联合训练场景中，实现了数据集用量的减少（高达30%）。同时，在大规模的单语、多语、跨语言检索基准测试中，tDRO显著提升了不同尺寸(500M、1.8B、4B、7B、8B)和不同基座(Qwen1.5、LLaMA3、Mistral-0.1)的大模型稠密检索（LLM-DR）的性能表现。

研究方法

问题定义

假设大模型稠密检索(LLM-DR)的训练集 Dtrain 包含 k 个数据集，每个数据集按照权重 α 的比例进行随机均匀采样 U，则本联合训练的数据分布定义如下:

面向大模型稠密检索(LLM-DR)的分布鲁棒优化 (DRO) 问题，旨在找到一组最优的 α，以最大化模型在下游任务上的性能。

研究挑战

在针对大模型稠密检索(LLM-DR)进行微调时，分布鲁棒优化(DRO) 面临两大挑战：

挑战 1: 采样策略冲突

经典的分布鲁棒优化(如 GroupDRO)算法，在一个批次(Batch)内读取并对比来自所有领域的数据。然而，稠密检索 (LLM-DR) 微调基于对比学习原理，为保证批次内负样本 (In-batch negatives)与跨批次负样本(Cross-batch negatives)的质量，一个批次(Batch) 内仅能读取来自同一领域的数据，否则将会影响最终性能表现。

挑战 2: 领域损失不可比

分布鲁棒优化（GroupDRO）等算法在鲁棒优化中，动态对比不同领域的损失(Loss)大小，损失越高，权重越大。然而，这种方法并不适用于大模型稠密检索(LLM-DR)中的异质数据。以下表中三个数据集的损失大小为例，在 Qwen1.5-0.5B模型上，Yahoo answers 的损失是 MS-MARCO 的 3 倍，DuReader 的 5 倍。如果直接以领域损失作为权重分配的依据，那么 Yahoo answers 将始终会得到更高的权重，这种优化方式显然是不合理的。

Table 1: Yahoo answers (Title-Answer)，MS-MARCO 和 DuReader 数据集损失大小对比。LLM-DR 模型以 Qwen1.5-0.5B 为基座，均匀采样训练了 1k steps。

tDRO：任务级分布鲁棒优化

为了解决上述两个挑战，我们创新性地设计了任务级分布鲁棒优化(Task-level Distributionally Robust Optimization, tDRO)算法。整体优化流程如下图所示:

为了解决 LLM-DR 和 DRO 之间的采样冲突，tDRO算法将分布鲁棒优化(DRO)与稠密检索微调(LLM-DR Fine-tuning)两个过程进行分离。首先，利用代理模型 (Proxy Model)只使用困难负样本(Hard Negatives, HN)来学习鲁棒的数据分布。随后，将数据分布迁移到 LLM-DR 中，结合三种负样本（困难负样本、批次内负样本、跨批次负样本）进行完整训练。

针对领域损失不可比的问题，tDRO 算法额外引入了 一个已训练好的参考模型(Reference Model)，通过“代理损失 / 参考损失”的比值作为领域损失度量 M(Group Loss Measurement)，来代表对应领域的进步空间。具体的损失定义、优化目标、 权重更新公式如下:

InfoNCE Loss

在对比学习的 step t 中，代理模型以 1/k 的比例采样所有数据，并为每一条数据生成一组 [查询 q，正例文档 d+，困难负例文档 d−hard]。代理模型使用 InfoNCE loss 作为损失函数，其计算过程如下:

其中，代理模型初始化自 Qwen1.5-0.5B，τ 是对比学习的温度系数，实验中设置为 0.002。

优化目标

tDRO 将每个数据集视为一个领域(Group)，并通过双优化目标交替更新领域权重 α 与代理模型参数 θ，以降低表现最差领域损失度量(Worst Group Loss Measurement) 的上确界:

权重 α 更新。tDRO 通过指数梯度上升更新权重 α，领域 g 的损失度量 M 越大，权重更新量越大。

*其中 ηα 为权重的学习率。

相对损失度量。如何获取一个可比的、有代表性的损失度量是 tDRO 优化的关键之处。tDRO 以平均采样策略额外训练了一个参考模型(Reference Model)。通过代理损失除以参考损失得到相对损失度量 M，以代表模型在该领域的进步空间。M 越大，代表进步空间越大，其权重更新量也越大。

代理模型更新。权重 α 更新后，代理模型以 α 作为系数计算 InfoNCE Loss，并利用 AdamW 优化器进行梯度下降更新。

LLM-DR 微调

在 tDRO 阶段学习到鲁棒的权重分布后，这些权重将按以下两种方式迁移至 LLM-DR 微调过程:

• Top-rated dataset selection：直接选取权重较大的前几组，以均匀采样策略进行训练。实验中我们发现，选取前 70% 的数据集性能最优。
• Sample Ratio Reweighting：直接按照 tDRO 的权重分布采样所有数据集。

LLM-DR 微调同样使用了 InfoNCE Loss，并结合了困难负样本(Hard Negatives)、批次内负样本(In-batch negatives)、跨批次负样本(Cross-batch negatives)三种负样本进行训练。需要注意的是，LLM-DR 在同一个批次(Batch)只采样一个来源的负样本，以保证负样本的质量。

实验结果

实验设置

本文实验概括如下表所示，在LLM-DR 训练的 base model 中，我们选择了 Qwen-0.5B、1.8B、4B、7B，Mistral-7B 和 LLaMA3-8B。

检索性能

本次对检索基线性能与 tDRO 优化后的鲁棒权重进行了实验。实验结果显示，tDRO 优化在多语言检索 MIRACL、跨语言检索 MKQA、单语言 BeIR 三个检索基线上带来了显著的性能增益。同时，我们将 tDRO 和其他领域损失度量所得权重分布进行了对比，发现 tDRO 将大约 9 个(占总体数量的 30%)单语数据集的权重降低至 0.01 以下。

进一步地，在 Dataset Selection Top-70% 实验组中，我们尝试将它们从训练中全部去掉，所有的 LLM-DR 的多语言、跨语 言、单语言检索性能均得到了提升，这证明了 tDRO 方法的有效性。

此外，我们还对比了不同领域损失度量优化后的权重变化。如上图所示，无论是直接使用领域损失，还是使用“领域损失-参考损失”，两者均明显偏向于 Yahoo 数据集；而 tDRO 通过“相对领域损失度量”，在不同任务之间，得到了更加合理的进步空间。下表的基线测试也证明，tDRO 的性能增益最好，而其他设计则稍显逊色。

实际应用与价值

优化大模型应用的检索增强框架

在大模型落地应用中，检索增强框架发挥着举足轻重的作用，不仅极大地提升了大模型的知识获取能力和问题回答的事实准确性，还有效降低了大模型的幻觉现象。澜舟科技始终致力于满足澜舟智库、澜舟智会和澜舟智搭系列产品的独特需求特点，持续优化检索增强相关技术，并成功设计和应用了多跳RAG、大跨度RAG、RAG指令微调等多种技术。在此基础上，我们进一步推出了tDRO技术，这一创新成果不仅实现了检索模型的大模型化，还进一步提高了检索模型的效果，使得检索增强框架在这些产品中发挥更加出色的作用，从而为客户带来更好的产品使用体验。

支撑跨语言、跨领域的应用场景

随着全球化和多领域知识的快速，大模型应用系统正面临着跨语言和跨领域应用的巨大挑战。本研究提出的技术框架为这些挑战提供了坚实的技术支撑。基于该框架训练的检索模型具备强大的跨语言和跨领域能力，能够在不同语言环境中实现精准检索，助力国际化企业跨越语言障碍，提高用户体验。同时，通过整合金融、医疗、法律等多个领域的训练数据，该框架显著增强了系统的多领域适用性，为企业跨部门的知识管理和协作提供了有力支持。此外。该技术还能整合多种数据来源，促进跨领域知识共享，帮助企业构建统一的知识中心，实现信息的全面覆盖与高效获取。

总结

在这篇文章中，我们设计了一种面向大模型稠密检索的任务级分布鲁棒优化(Task-level Distributionally Robust Optimization, tDRO)算法。该算法通过学习不同领域任务的鲁棒数据分布，成果减少了30%的数据集用量。在多种不同基座与不同大小参数量的大模型稠密检索微调中，tDRO算法均展现出了显著的检索性能提升。这一成果不仅证明了tDRO算法的有效性和实用性，也为大模型应用的检索增强框架的优化提供了新的思路和方法。

代码开源地址：https://github.com/tdro-llm/tdro

模型、数据开源地址：https://huggingface.co/tdro-llm

完整论文地址：https://arxiv.org/pdf/2408.10613