How to Train Your DRAGON: Diverse Augmentation Towards Generalizable Dense Retrieval
问题 (Problem)
传统的稠密检索(Dense Retrieval, DR)模型普遍存在一个核心问题:在监督式评测和零样本(Zero-shot)评测之间存在明显的性能权衡(trade-off)。具体来说,一个在特定数据集(如 MS MARCO)上通过监督学习训练得很好的模型,在从未见过的新领域(零样本场景)上往往表现不佳,反之亦然。
如 Figure 1 所示,大多数现有的检索器(除了参数量巨大(4.8B)的 GTR-XXL)都分布在一条倾斜的直线上,显示了“监督式评测”得分(X轴)和“零样本评测”得分(Y轴)之间的负相关关系。当时的普遍观点认为,要打破这种权衡,必须大幅增加模型(如BERT-base)的容量。
本文的核心目标:挑战上述观点,证明在不增加模型大小(仍使用 BERT-base 级别)的前提下,通过一种更优的训练方法,可以训练出一个在监督式和零样本场景下均达到顶尖水平(State-of-the-Art, SOTA)的通用稠密检索器。
Figure 1 解读:

- 坐标轴:
- X轴: Supervised evaluation (MS MARCO Dev: RR@10) - 衡量模型在见过的、有监督数据上的性能。
- Y轴: Zero-shot evaluation (BEIR 13 subsets: nDCG@10) - 衡量模型在未见过的新领域的泛化能力。
- 数据点: 图中的每个点代表一个检索模型。不同形状代表不同类型的检索器(如圆形代表稠密检索,方形代表稀疏检索)。
- 蓝色趋势线: 清晰地揭示了大多数模型存在的性能权衡。
- DRAGON (红色点): 本文提出的模型,显著地位于趋势线的右上方,表明它同时在监督和零样本任务上取得了高分,成功打破了这种权衡。
方法 (Method)
作者提出了一个名为 DRAGON (Dense Retriever trained with diverse AuGmentatiON) 的模型。其核心思想是构建一个统一的数据增强(Data Augmentation, DA)框架,通过多样化的查询增强和渐进式的相关性标签增强来训练模型。
统一的数据增强框架
作者将许多现有的稠密检索改进方法(如知识蒸馏、对比预训练、伪查询生成)都归纳到数据增强的框架下,主要分为两类:
- 查询增强 (Query Augmentation): 增加训练查询的数量和多样性。
- 相关性标签增强 (Relevance Label Augmentation): 为每个查询提供更丰富、更准确的正负样本标签。
查询增强
DRAGON 探索了两种主要的查询增强方式:
- 句子裁剪 (Sentence Cropping): 一种计算成本低廉的方法,直接从语料库的文档中裁剪出句子作为查询。这种方法可以轻松地将查询规模扩大到千万级别。
- 伪查询生成 (Pseudo Query Generation, GenQ): 使用大型语言模型(如 T5)为文档生成类似人类提问的查询。这种方法生成的查询质量高,但计算成本昂贵。
最终,DRAGON 采用混合方式,将一半的裁剪句子和一半的生成式查询混合作为训练数据,以兼顾数量、多样性和质量。
多样化和渐进式的标签增强
这是 DRAGON 方法的核心创新点。传统的标签增强通常依赖于单个、强大的“教师模型”(如交叉编码器 Cross-Encoder, CE),但作者认为单个教师的视角是有限的。
核心假设:不同的检索器(如稀疏、稠密、多向量模型)捕捉到的文本相关性信号是不同的(例如,词汇匹配 vs. 语义匹配)。为了让学生模型学到更通用的匹配能力,应该从多个具有不同优势的教师那里学习。
具体实现:渐进式监督 (Progressive Supervision)
DRAGON 不会一次性将所有教师的知识都灌输给模型,而是采用一种类似课程学习(Curriculum Learning)的策略,分阶段、逐步地引入新的教师。
- 教师团队: DRAGON 使用了一组多样化的教师模型,包括:
- 稀疏检索: uniCOIL, SPLADE++
- 稠密检索: Contriever, GTR-XXL
- 多向量检索: ColBERTv2
- 训练过程: 训练分为多个迭代(Iteration)。
- 第一阶段: 只使用第一个教师(如 uniCOIL)生成的排序列表来为增强查询构建训练样本(正负例)。
- 第二阶段: 引入第二个教师(如 Contriever),此时,对于每个查询,模型会从第一和第二个教师的排序列表中均匀采样一个作为监督信号。
- 后续阶段: 逐步引入更多的教师,在第 个阶段,模型会从前 个教师的监督信号中均匀采样。
Figure 2 图解:

该图直观地展示了渐进式监督的过程。
- Iteration 1: 只有红色类别的标签(来自教师1)。模型从这些正负样本中学习。
- Iteration 2: 在原有基础上,增加了绿色类别的标签(来自教师2)。此时训练数据的“视野”扩大了,模型需要学习兼容两种不同的相关性判断。
- Iteration 3: 进一步增加了蓝色类别的标签(来自教师3)。模型学习的监督信号变得更加丰富和复杂。
这个过程引导模型从简单、单一的视角逐步过渡到复杂、多元的视角,从而更有效地学习通用性。
数学公式
DRAGON 的训练遵循标准的稠密检索对比学习框架。
- 相似度分数: 查询 和文档 之间的相似度通过它们向量表示的点积计算:
其中 和 是查询和文档的编码器(通常是 BERT),输出的是 [CLS] token 的向量表示。
- 损失函数: 使用 InfoNCE 损失函数。对于一个查询 、其相关的正样本 和一组负样本 ,目标是最小化以下损失:
这个损失函数旨在拉近查询与正样本的距离,同时推远与负样本的距离。在 DRAGON 的训练中, 和 是从教师模型生成的排序列表中采样得到的。
Baseline
DRAGON 与多种先进的检索模型进行了比较,涵盖了不同的架构和训练技术:
- 不同架构模型:
- SPLADE++: 强大的稀疏检索模型。
- ColBERTv2: 高效的轻量级晚期交互(multi-vector)模型。
- GTR-XXL: 参数量巨大的稠密检索模型(4.8B)。
- 基线稠密检索模型 (BERT-base):
- 知识蒸馏: RocketQAv2, CL-DRD。
- 对比预训练: coCondenser, Contriever, COCO-DR。
- 掩码自编码预训练: COT-MAE, RetroMAE。
- 无监督领域自适应: GPL, PTR。
数据集 (Datasets)
- 训练数据来源:
- MS MARCO Passage Corpus (8.8M passages): 用于进行查询增强和标签增强,是生成所有训练数据的唯一语料库。
- 评测数据集:
- 监督式评测:
- MS MARCO Dev: 标准的领域内(in-domain)评测集。
- TREC Deep Learning (DL) 2019 & 2020: 拥有更精细人工标注的监督评测集。
- 零样本评测:
- BEIR: 一个包含18个不同领域、不同任务的异构零样本评测基准。
- LoTTE: 包含多个 StackExchange 社区问答数据的零样本评测集。
- 监督式评测:
可复现性 (Reproducibility)
- 代码: 作者在论文中提供了 GitHub 仓库链接:
https://github.com/facebookresearch/dpr-scale,包含了代码和模型 checkpoints。 - 算力: 训练成本非常高昂。论文在“局限性”部分明确指出,训练一个完整的 DRAGON 模型需要:
- 32 块 A100 (40 GB) GPUs
- 耗时 5 天
- 此外,还需要预先准备好所有训练好的教师模型,并使用它们对千万级别的增强查询进行推理以生成标签,这本身也是一个巨大的计算开销。
可改进的几个点 (Potential Improvements)
论文自身也指出了几个未来可以探索的方向:
- 降低训练成本: 当前的训练数据规模高达 2800 万,其中可能包含大量重复或低质量的查询。设计一种有效的数据筛选或去重方法,可以在不牺牲性能的前提下大幅降低训练成本。
- 与领域自适应结合: DRAGON 已经具备了很强的零样本能力。如果将其作为基础模型,再结合 GPL、PTR 等领域自适应技术进行微调,可能会在特定目标领域上取得更好的性能。
- 与其他预训练方法结合: DRAGON+ 的实验表明,掩码自编码预训练(如 RetroMAE)与本文的对比学习方法是互补的。未来可以进一步探索如何更好地融合生成式和对比式预训练方法。
- 特定领域的知识增强: 在 LoTTE 数据集的科技主题上,DRAGON 的表现不如 SPLADE++ 和 ColBERTv2。这表明,在训练中引入特定领域的语料库(如科学文献)可能会弥补这一不足。
可以被引用的结论 (Key Takeaways)
- 成功打破性能权衡: 本文首次证明,一个标准的 BERT-base 大小的稠密检索器,无需增加模型参数,就能同时在监督式和零样本评测中达到 SOTA 水平,打破了两者之间的固有权衡。
- 多样化监督优于单一强监督: 与普遍认为“应使用最强的交叉编码器作为教师”的观点相反,本文发现,使用来自多个不同类型检索器(稀疏、稠密、多向量)的多样化监督信号,对于训练模型的泛化能力更为关键。
- 渐进式学习策略的有效性: 渐进式监督(Progressive Supervision)作为一种课程学习策略,能有效引导模型学习复杂且多样的相关性信号,其效果优于一次性混合或融合所有监督信号。
- 廉价数据增强的意外价值: 实验惊人地发现,通过句子裁剪这种廉价方式生成的增强查询,在提升模型泛化能力方面,效果甚至优于使用昂贵语言模型生成的“类人”查询(GenQ)。这为未来低成本、大规模训练通用检索器提供了新的思路。
- 统一的数据增强框架: 论文提出的数据增强(DA)框架,为理解和归类稠密检索领域的各种训练技巧(如知识蒸馏、对比学习等)提供了一个统一且有价值的视角。