ANCE:用全语料 ANN 难负样本训练 dense retriever

ANCE(Approximate Nearest Neighbor Negative Contrastive Learning for Dense Text Retrieval,2020,Microsoft,ICLR 2021)是一篇 2020 年的老文章,但 dense retrieval 训练里 hard negative mining 的标准做法基本是从这里定型的,DPR 系列、E5、BGE 都沿用了这套框架。其核心思路:在训练过程中持续维护一份 ANN 索引,从整个语料里挑当前模型最难区分的负样本去训 dense retriever,索引以异步方式定期刷新

论文有两层贡献。一是用 importance sampling 的方差分析指出,dense retrieval 训练时常用的 in-batch 与 BM25 负样本梯度范数接近零,是收敛瓶颈。二是给出全语料 ANN 难负样本采样的具体方案,并解决"索引必须随模型同步更新"这个工程瓶颈。在 TREC 2019 DL、MS MARCO、NQ、TQA 上,BERT-Siamese 配 ANCE 训练后,文档检索 NDCG@10 达 0.628(MaxP),passage MRR@10 达 0.330,Top-20 Coverage 在 NQ 和 TQA 上均超过 DPR,生产环境 8B 语料的离线检索质量相对提升约 14%~15%。

论文地址: arxiv.org/abs/2007.00808

问题:dense retrieval 为什么打不过 BM25

dense retrieval(DR)把 query 与 document 编码成稠密向量,靠 ANN 在向量空间检索。理论上 DR 能学到词面之外的语义匹配,但 2020 年前后的现实是,端到端学的 DR 在文档检索上常常打不过 BM25。

论文把 gap 归到负样本选择上。第一阶段检索的特殊之处在于,“不相关文档"指的是整个语料 $C$ 减去相关集合 $D^+$,规模百万到十亿,一次只能采一小撮。常见做法有三类:BM25 top 文档当负样本、in-batch 用其他 query 的正样本当负样本、或两者混合。这些做法在 DR 上的实测效果都很有限,论文要回答的是为什么。

梯度方差视角下的局部负样本失效

作者从 importance sampling 写出 SGD 一步的收敛速率:

$$E_t = 2 E_{P_{D^-}}(g_{d^-})^T(\theta_t - \theta_*) - 2 E_{P_{D^-}}(g_{d^-})^T E_{P_{D^-}}(g_{d^-}) - 2 \text{Tr}(V_{P_{D^-}}(g_{d^-}))$$

其中 $g_{d^-}$ 是负样本 $d^-$ 上的加权梯度。关键在最后一项:减去梯度估计的方差。要让方差小、收敛快,最优采样分布是

$$p_{d^-} \propto |\nabla_{\theta_t} l(d^+, d^-)|_2$$

按梯度范数大小成比例采样。梯度范数大的负样本对 loss 推动大,是信息量大的样本。

对常用的 BCE 与 pairwise hinge loss 可以验证,$l(d^+, d^-) \to 0$ 时梯度范数也 $\to 0$,loss 已接近零的负样本对训练几乎没贡献。记 $D^-$ 为真正难区分的负样本集合,batch size 为 $b$,DR 训练上有两条经验事实:$b \ll |C|$ 且 $|D^-| \ll |C|$。两者相乘,一个随机 mini-batch 包含至少一个有信息量负样本的概率近似 $\frac{b \cdot |D^-|}{|C|^2}$,趋近于零。结论很清楚:DR 训练空间太大,难负样本太稀,光在 batch 里找根本碰不上。

全语料 ANN 难负样本与异步索引刷新

ANCE 的训练目标是

$$\theta^* = \arg\min_{\theta} \sum_q \sum_{d^+ \in D^+} \sum_{d^- \in D^-_{\text{ANCE}}} l(f(q, d^+), f(q, d^-))$$

$D^-_{\text{ANCE}}$ 取自当前 DR 模型在整个语料上的 ANN top 检索结果,去掉真实正样本。按上一节的分析,这些就是当前模型梯度范数最大、对训练贡献最高的样本。

工程难点是,模型每个 mini-batch 都在更新,理论上 ANN 索引也要跟着更新,但每步重新编码全语料加重建索引完全不可行。ANCE 借鉴 REALM 的做法:两个并行进程,Trainer 持续训练,Inferencer 拿最近 checkpoint $f_k$ 把全语料重新编码并重建 ANN。Trainer 在等的过程中继续用旧索引 $\text{ANN}_{f_{k-1}}$ 采的负样本训练,每过 $m$ 个 batch 才换上新索引。

ANCE 异步训练:Trainer 用旧索引继续训,Inferencer 在另一组 GPU 上重新编码整个语料并刷新 ANN 索引

代价是 Trainer 用的负样本始终略微"过期”。论文用 1:1 GPU 切分(4 卡训 + 4 卡推),实测每 10k batch 刷一次索引这一程度的过期对训练影响很小。

实现上沿用 BERT-Siamese 双塔,dot product 相似度,NLL loss。长文档两种变体:FirstP 取前 512 token,MaxP 切成最多 4 个 512-token 段、段级分数取 max。所有 DR 模型先用 BM25 检索出的难负样本做 warm-up 再切到 ANCE,否则冷启动的 ANN 接近随机,采不到难负样本。

实验结果

TREC 2019 DL Track 的文档检索 NDCG@10:BM25=0.519,最强 sparse 方法 DeepCT(用 BERT 学 token 重要性后灌进倒排索引)=0.554,混合负样本 DR baseline BM25+Rand=0.557,ANCE FirstP=0.615,MaxP=0.628。论文强调的关键观察是,在所有 DR 负样本构造方法里,只有 ANCE 让 BERT-Siamese 在文档检索上稳定超过 sparse 方法。

passage 检索 MRR@10:BM25=0.240、BM25 Neg=0.299、DPR 风格的 BM25+Rand=0.311、ANCE FirstP=0.330。

OpenQA 上 NQ 与 TQA 的 single-task Top-20 Coverage:ANCE 分别 81.9 和 80.3,DPR 是 78.4 和 79.4。接 reader(NQ 用 RAG-Token,TQA 用 DPR Reader)后端到端答案准确率 NQ 从 41.5 提到 46.0,TQA 从 56.8 提到 57.5。

生产搜索环境:250M 语料、768 维、精确 KNN 相对提升约 18.4%;8B 语料、64 维下 KNN 提升 14.2%、ANN 提升 15.5%。

效率:ANCE 离线推理(编码全语料 + 建 ANN)约 10h;在线检索每 query 加 100 文档只要 11.6ms,对比 BM25+BERT 重排的 1.42s 约 100x 加速。

表示空间一致性与收敛行为

最值得信的证据是 reranking 与 retrieval 之间的 gap 被抹平。同样的模型既可以当做 reranker(对 BM25 top-100 重排),也可以当做 retriever(对全语料 ANN 检索),两种用法分数本来差得很大。ANCE 把两者拉到几乎相同(FirstP doc:rerank 0.641 vs retrieval 0.615;MaxP:0.671 vs 0.628),其它 DR baseline 两列差距都更大。换句话说,ANCE 训出的表示空间在 reranking 时找得到的相关文档,在全语料 ANN 检索时也找得到。论文借此说"interaction-based BERT 不是必需的",这是对当时神经 IR 主流认知的一处修正。

收敛分析上,ANCE 的训练 loss 长期保持在高位,NCE Neg 和 Rand Neg 这类局部负样本的 loss 早早降到接近零。梯度范数上 ANCE 比这两种局部负样本大几个数量级,与上一节的理论预测一致。论文还测了局部负样本与"测试时真正难的负样本"的重叠率:NCE Neg 与 Rand Neg 都是 0%,BM25 Neg 是 15%,ANCE 起点 63%,收敛时按定义到 100%。

刷新频率、warm-up 与 sparse 召回差异

异步索引刷新频率是关键超参。论文设置是每 10k batch 刷一次,1:1 GPU 切分;5k 比较平稳,20k 在高学习率下出现波动,但合适 LR 下都能跑。GPU 切分过度偏向 Trainer 会让索引明显过期。

BM25 warm-up 是必需的。冷启动 ANN 接近随机,采不到难负样本。这意味着 ANCE 不是独立训练范式,而是 sparse 负样本预热 + dense 难负样本精调的两阶段流程。

DR 与 sparse 检索的结果重叠比想象中小。ANCE 与 BM25 的 Top-100 重叠只有 17%~25%,说明 dense retrieval 学到的不是 BM25 的近似,而是另一套相似度度量。论文给出的 case study 显示,ANCE 召回的 irrelevant 文档往往"语义相关但不完全切题",与 sparse 检索的失败模式不同。

小结

ANCE 做的事可以拆成两步看。一是诊断:用 importance sampling 的方差分析证明 in-batch 与 BM25 局部负样本在 DR 训练里梯度范数趋零,是收敛瓶颈而非工程细节。二是方案:从全语料用异步刷新的 ANN 索引采难负样本,配 BM25 warm-up 解决冷启动。