BlaXuan 发表于 2022-1-16 17:08

Optimizing Dense Retrieval Model Training with Hard ...

Source:

Code:

概览:

搜索结果排序一直是信息检索领域核心的研究问题。传统的检索模型主要依靠词级别的匹配信息,近年来一些研究者使用基于稠密向量的检索模型(DR,Dense Retrieval)来进一步提升检索性能。训练一个好的稠密向量检索模型依赖于训练数据的采样。
本文针对如何有效的采样用以训练稠密向量检索模型负样本进行了研究。作者首先在理论上对现有的训练方式进行了深入分析,并尝试解释为什么在训练中采样困难负样本(hard negatives)比随机采样排序效果更好。通过分析发现,目前大多数研究都使用随机采样(random sampling)或者静态困难负样本采样(static hard nagative sampling)。其中,随机采样往往优化模型整体排序误差,对检索top-k文档来说是次优选择。而静态困难负样本采样虽然以top-k文档排序误差为优化目标,但由于负样本会随着模型训练对应排序降低,会给训练带来潜在的风险。
因此,作者在本文中提出了两种新的训练策略:
1、稠密检索的稳定训练算法(STAR, Stable Training Algorithm for dense Retrieval):STAR策略把静态的困难负样本采样和随机采样结合,在优化top-k排序指标的同时保证训练过程稳定。
2、直接优化排序性能的查询端训练算法(ADORE, Algorithm for Directly Optimizing Ranking pErformance):ADORE策略进一步用动态困难负例来训练查询编码器,提升检索性能。
在公开数据集上的实验表明,这两种训练策略能有效地提升现有稠密向量检索模型的检索性能,并且两者的组合可以获得最佳的性能。同时,这两种训练策略的运行效率也显著优于现有的基线模型。
Related work:

DR模型的训练方法除了知识蒸馏外,还主要采用负采样方法。
有研究使用随机负采样,使用随机负样本近似召回的任务。使用mini-batch训练,将同一批次中其他查询的相关文档作为负样本,并发现增加小批量随机负样本的数量是有益的。
有研究使用困难负样本方法,使用BM25top文档作为困难负样本。使用预热DR模型在训练中预先检索top文档作为困难负样本,其定期重建索引、刷新困难负样本,大大增加了计算成本。
静态困难负样本采样:困难负样本在整个训练过程中是固定的。然而有研究发现,静态困难负样本采样不一定有利于排序,甚至恶化了排名性能。
任务:

学习目标:DR模型通常采用pairwise loss进行训练,由于训练过程中不可能训练全部样本,因此不同的采样策略看成是为负样本文档设置不同的权重 https://www.zhihu.com/equation?tex=w%EF%BC%88d%5E%7B-%7D%EF%BC%89


其中:


描述了正样本文档的排序位置π()与训练损失间的关系。其中δ()是排名比更高的相关文档的数量。
理论分析:

1、Random VS. Hard Negatives

(困难负样本到底是怎样帮助优化检索模型的?)

[*]Random Negative Sampling:
在整个语料库中统一地进行负样本采样,相当于使用随机负样本最小化π()。


该优化目标存在一个关键问题:关于单个查询的损失函数是无限的,π()与整个语料库的大小一致。因此,整体pairwise loss主要学习语料库之间的差异,模型很难专注于提高top-k文档排序的性能。

[*]Hard Negative Sampling:
从所有负样本中找到top-k个负样本进行训练,优化目标:


K:困难负样本的数量   https://www.zhihu.com/equation?tex=K_%7Bq%7D :topK个负样本的排序位置
因此困难负采样方法是优化top-k的pairwise loss,而不是整体的pairwise loss,单个查询的损失受限于K,相比随机负采样方法可以更好地提高top-k文档排序的性能,对检索结果的优化更鲁棒。


说明如果DR模型能够将所有相关文档排在较高的位置,那么在训练中使用随机负采样和困难负采样将得到相同的最优参数。
但是对于需要关键字匹配的查询,DR模型可能会失效,这种情况下,随机采样策略会过度关注于困难查询,但该问题在主流的评估指标中并不会被捕捉到。
2、Static VS. Dynamic Hard Negatives


[*]Static Hard Negatives:


其中是预先检索出的困难负样本集合,使用最高排序位置定义负样本质量Φ()


可以看出Φ()非常松散,理论上取值可以非常大,使用MRR评价排序性能,MRR下限值为:


训练损失为0的“理想”DR模型,结果为 https://www.zhihu.com/equation?tex=%5Cfrac%7B1%7D%7B%7CC%7C%7D ,该结果是我们不能接受的。
因此,静态困难负采样在训练过程中是存在风险的。

[*]Dynamic Hard Negatives:
在训练过程中,根据DR模型参数的变化,同步提供排名最高的困难负样本。


最高排序质量Φ()为负样本质量:


MRR下限值为:


因此,在理论推导上来讲,动态困难负采样优于静态困难负采样。
算法设计:

作者提出两种改进的训练策略——STAR & ADORE,来更好地训练DR模型。
STAR算法:

由于静态困难负采样是不稳定的,因此作者提出一种稳定的密集检索训练算法Stable Training Algorithm for dense Retrieval(STAR)。
1、使用静态困难负样本来近似真实的动态困难负样本
使用一个预热模型检索所有训练查询的top文档,将其作为静态困难负样本,这些负样本在训练过程中不会被更新。
2、为稳定训练,引入随机采样因素
优化目标:



https://www.zhihu.com/equation?tex=L_%7Br%7D%EF%BC%88%5Ctheta%EF%BC%89 为方程5定义的随机负采样的损失函数,为方程7定义的静态困难负采样的损失函数,α为超参数。
如果静态困难负样本可以很好地接近动态困难负样本,值比较大,主导训练过程。如果不能接近于动态困难负采样,STAR算法近似退化为一种随机负采样方法,该方法仍比最差的静态困难负采样好得多。
3、最后使用重用策略提高训练效率
在训练时没有再次输入随机负样本,而是重用同一批中其他样本近似随机负样本。如图1,重用策略可视为in-batch训练策略的改进版本


4、损失函数采用RankNet pairwise loss


ADORE算法:

一种直接优化排名性能的查询端训练算法
a query-side training Algorithm for Directly Optimizing Ranking pErformance (ADORE)


1、使用动态困难负样本
在训练之前,ADORE使用一个预先训练的文档编码器预先计算文档embedding,构建文档索引,其在整个训练过程中是固定的。每次训练迭代中,对一批查询进行编码,使用查询向量检索相关的top文档,作为动态困难负样本。ADORE利用这些负样本来训练DR模型。这是第一个使用动态困难负样本的DR训练方法,作者已通过理论分析证实其有效性。
2、直接优化IR指标
ADORE在每一步都执行检索,因此可以利用listwise方法直接优化排序性能。作者使用LambdaLoss更好地优化IR指标。


ΔM是交换正负样本的排序位置给IR指标带来的变化大小,研究表明,通过在RankNet loss上增加ΔM,该方法可以直接优化IR指标。并且有研究认为,当训练集足够大时,该方法会根据相应的指标生成最优模型。
3、End-to-end Training
不同的文档索引压缩方法会导致不同的最优DR参数,这一点在之前的研究中多被忽略。ADORE可以利用推理中使用的实际文档索引执行端到端的训练。ADORE算法缓解了训练与推理的差异,可以获得更好的排序能力。之后作者也进行了实验验证。
Combine STAR and ADORE:

ADORE算法和STAR算法有着各自的优势,ADORE算法可以直接优化排序性能,而STAR算法不能;STAR算法同时优化了查询编码器和文档编码器,而ADORE算法只优化查询编码器。本文作者结合两个策略,采用STAR算法对文档编码器进行预训练,然后使用ADORE算法对查询编码器进一步训练。
Experimental Settings:

1、数据集:TREC2019Deep Learning(DL)Track
该任务侧重于ad-hoc检索(一种特定的检索),包括段落检索和文档检索任务。
2、评估指标:MRR@10、NDCG@10、R@100
3、Baselines:
Sparse Retrieval:BM25、DeepCT
Dense Retrieval:

[*]随机负采样baseline:Rand Neg从整个语料库中随机抽取负样本,In-Batch Neg使用同一批次中其他查询的相关文档作为负样本。
[*]静态困难负采样baseline:BM25 Neg使用BM25最优候选文档作为负样本,ANCE采用一个预热模型检索静态困难负样本。
[*]知识蒸馏baseline:TCT-ColBERT使用了ColBERT作为teacher model。
Cascade IR:本文还采用级联系统进一步比较,包含以BM25为一级检索器的best LeToR模型和BERT模型。
4、实现细节:
所有的DR模型使用RoBERTa模型作为编码器,使用内积计算相似度,采用Faiss库进行高效的相似度搜索。对于段落检索和文档检索任务,文档分别被分割为最多120个、512个token,前200名的文档作为负样本。DR baselines损失函数使用RankNet loss,训练后的模型直接在文档任务上进行评估以产生最好的检索性能。STAR采用BM25 Neg模型作为预热模型,与ANCE相同,因此其结果可直接比较。
STAR使用Lamb优化器,ADORE使用AdamW优化器。公式15中的M在段落检索和文档检索任务中分别为MRR@200和MRR@10。注意到ADORE可以进一步训练其查询编码器来改进训练后的DR模型,比如ADORE+Rand Neg表示其使用Rand Neg训练过的文档编码器,进一步训练查询编码器。
5、实验结果:
作者通过系列实验验证其理论分析和提出算法的有效性,主要关注以下三个研究问题:
RQ1:随机负样本和困难负样本是如何影响优化目标的?
RQ2:静态困难负采样方法在实践中表现如何?
RQ3:作者提出的训练方法有效性和效率如何?
1、之前的理论表明,随机负采样面临一个问题:一些困难查询可能会主导整个训练过程。为验证是否真的存在该现象,作者在图3中绘制了MARCO Dev Passage数据上的Rand Neg模型每个查询的pairwise损失的分布,结果显示0.2%的困难查询在整体的pairwise loss中贡献了60%,该问题在实践中是非常严重的。


为此作者评估了不同模型的总体的pairwise errors和top-K(200)pairwise errors。


可以看出,随机负采样可以优化总体的pairwise loss,但不能最小化top-K pairwise loss。困难负采样方法STAR很好地优化了top-K pairwise loss。静态困难负采样BM25、ANCE取得了较差的top-ranking性能。ADORE有效提高了每个模型的top-ranking性能,其中使用STAR训练的文档编码器表现最好。因此,困难负采样可以更好地优化排名性能。
此外,该实验结果与之前的研究结果相悖,有研究者认为随机和困难负采样有相同的优化目标,困难负采样的优势在于收敛速度更快,但是我们在表1可以看出,随机负采样收敛性更好,整体pairwise loss优于困难负采样。图4b显示,随机负采样在初始化较好时会降低top-ranking性能。因此,随机和困难负采样是对不同的目标进行优化,而困难负采样的优势与收敛速度无关。


2、作者认为DR模型可能会很快地将静态困难负样本排序在非常低的位置,因此其无法近似动态困难负样本,同时静态困难负采样不一定能够优化排序性能。
图4a显示了静态困难负样本和真实动态困难负样本间的重叠率,可以看出重叠率随训练的增加而减小。因为静态困难负样本不能考虑到训练过程中dropout引入的随机噪声,重叠率总是小于70%,因此,动态困难负样本不能完全被静态的所取代。
图4b显示了排序性能,静态负采样过程非常不稳定并且很快就被DR模型排到了很低的位置,其出现了周期性的剧烈波动,大多数情况下其表现不如随机负采样。相反,动态困难负采样能稳定地提高排序性能。
虽然有研究通过调参很好地解决了以上问题,但作者接下来的实验表明,其提出的方法可以更好地优化排序性能,并获得较大的效率提升。
3、有效性


首先对于baselines的表现,我们可以看出随机负采样可以有效地训练DR模型,尤其表现在召回率上。静态困难负采样并不一定会比随机负采样带来性能的改进,其提高了top-ranking性能,但可能会破坏召回能力。
其次,可以看出STAR显著优于随机负采样,尤其是在top-ranking性能上。STAR还明显优于静态困难负采样方法,尤其是在召回率上。可见,STAR通过稳定训练过程,更好地优化了排序性能。
对于ADORE算法,其通过进一步训练查询编码器极大地提高了所有DR模型的性能。
将STAR与ADORE两者结合起来达到了最佳性能,显著优于所有的baselines,特别是有竞争力的知识蒸馏方法TCTColBERT和周期索引刷新方法ANCE。此外,该方法在文档检索任务上与BM25-BERT检索系统相当。
作者对ADORE进行了消融实验


图5结果可以看出动态困难负样本和Lambdaloss都极大地提高了排序性能。


图6中,ADORE使用BM25 Neg训练的文档编码器,并进一步训练查询编码器,训练后的ADORE将查询映射到更接近相关文档的位置,提升了检索性能。


不同的压缩索引,最优DR参数可能不同,为了验证ADORE通过端到端训练可以获得更好的性能,作者使用不同的压缩索引来训练和评估DR模型,如表3所示,可以看出端到端训练可以更好地优化不同压缩技术的排序性能。
4、效率
ANCE在效率上具有竞争优势,因此作者将其作为效率基线。
训练时间:


可见本文的方法收敛更快,一方面,在段落检索任务上,ANCE需要600k步,batch size为64,而ADORE需要60k步,batch size为32。另一方面,为定期更新静态困难负样本,ANCE对语料库进行迭代编码,每次耗时10.75小时,使用3个GPU。相反,STAR只编码一次,ADORE甚至没有这种开销。
计算资源:


ADORE可以通过压缩文档索引极大地节省计算资源。本文使用PQ【13】来压缩由BM25 Neg生成的文档向量,压缩PQ索引的内存使用情况和搜索质量如表5所示。
可以看出PQ显著降低了内存消耗并且性能损失较小。压缩后,ADORE可以在一个GPU上运行。PQ-96可以将GPU内存占用减少到3%,并且性能损失很小(0.326vs0.329)。更小的PQ进一步减小了内存占用。因此,ADORE更适用于内存受限的环境,但也要注意过度压缩(PQ=6)会严重降低检索质量,ADORE会退化为随机负采样(0.304vs0.301)。
创新点:

1、比较分析了稠密检索模型中不同的训练策略会导致不同的优化目标。从理论和实证两个方面论证了静态困难负样本采样策略有着降低样本排序的风险。
2、提出STAR算法和ADORE算法,将两者结合,使用动态困难负样本采样方法代替广泛使用的静态困难负样本采样方法,可以直接优化排序性能。同时引入随机采样也提高了稠密检索模型训练的稳定性。
问题:

虽然在本文中向量检索效果提升了不少,但是其实验数据集长度还是较短。在真实场景下,面对上万字的文档进行检索,仍然存在很大的挑战。因此,对于大规模文档中的稠密检索问题仍值得进一步的研究。
页: [1]
查看完整版本: Optimizing Dense Retrieval Model Training with Hard ...