基于树的端到端稠密检索模型

今天介绍的这篇文章由清华大学和华为联合发表，核心是提升向量检索的效果，在树检索的基础上，实现了索引构建和表示学习的端到端联合建模，提升了树检索的一致性。

论文标题：Constructing Tree-based Index for Efficient and Effective Dense Retrieval

下载地址：https://arxiv.org/pdf/2304.11943.pdf

机构：清华 & 华为

1、树检索

树检索是提升稠密向量检索效率的一种常用方法。在信息检索中，dense retrieval对于提升第一阶段检索效果有很大好处。Dense retrieval将query和document分别编码成向量，可以使用TwinBERT等类似的双塔模型拟合rank loss，得到query和document的向量。然后对于一个query向量，计算其和所有document向量的打分，选择topN的document作为第一阶段的检索结果。

然而，这种方式的计算开销很大，全库进行检索显然是不现实的。树检索就是为了提升这类dense retrieval而提出的一类算法。它的核心思路是，将所有document通过某种方式（如层次聚类）形成树结构，树中每个节点对应一个向量，代表以这个节点为根节点的子树。对于一个query向量，在树中做层次检索，每层只选打分topK的节点，进入到下一层匹配，下一层匹配只和上一层topK的节点的子节点进行匹配，按照这个逻辑递归进行（也即Beam Search，基于贪心的策略，每层只保留最相关的几个节点，逐层检索）。通过这种方式，可以显著降低检索空间，让dense retrieval在第一阶段的应用成为可能。

2、现有树检索的问题

现有的树检索模型，一般采用两阶段的方式：第一阶段训练query-document的双塔模型，拿到query和document的向量；第二阶段基于第一阶段训练好的向量，通过聚类算法构建层次树。整个过程如下图所示。

这种方式的弊端在于，两阶段的方式导致二者优化目标不一致，得到的并不是最优解。为了解决这个问题，本文提出了一种端到端的稠密向量学习+树索引构造的学习方式，实现了更高效的树稠密检索架构。

3、端到端训练

本文将Encoder训练和树学习融合到一起学习，新的流程如下图所示。由于cluster assignment是不可导的，优化比较困难，因此文中仍然使用Kmeans来做初始的聚类簇划分，重点将各个聚类节点的embedding通过对比学习的方式进行优化。

具体来说，文中首先利用原来的方法进行初始的聚类簇划分，然后使用对比学习优化每个cluster的embedding。对于一对query和document正样本，在树中采样它的正样本cluster节点和负样本cluster节点，然后做query和正负cluster节点的对比学习，以此优化cluster embedding，让cluster embedding能更好的进行beam search，尽量避免出现相关document在上层节点分支就被过滤掉的问题。构造正负样本的示意图如下。

4、带overlap的cluster

基于kmeans的聚类方法生成的树，每个document只会在一个叶子结点里。但是一个document可能对应不同的语义，对应不同的类别，这样的划分方法可能导致高质量document无法召回。因此，本文对cluster assignment也进行了优化，让一个document可以出现在多个叶子结点中。

核心做法如下图所示，利用query与叶子节点、document的关系矩阵进行融合，来实现overlap assignment。基于query-document的相关性矩阵，以及query-leaf node的assignment矩阵，可以获取到query与document、leaf node的相关关系。将这两个矩阵相乘，得到每个document属于每个lead node的关系，基于这个关系矩阵进行冗余document的挂载。其原理为，如果q1和l1、l2相关，且q1和d2、d4也相关，那么d2、d4都应该出现在l1、l2对应的叶子节点中，以提升树构建和相关性label的一致性。

5、实验结果

整体实验结果如下表所示，JTR方法带来了比较显著的效果提升。