SetRank: Learning a Permutation-Invariant Ranking Model for Information Retrieval
来自SIGIR20，代码已开源https://github.com/pl8787/SetRank
Introduction
传统的排序学习模型通常是基于概率排序原则Probability Ranking Principle （PRP）设计的，概率排序原则假设每个文档都有某种满足特定信息需求的独特的概率。因此，文档的排序分数是单独分配的，并且彼此独立。基于PRP的排序学习算法有固有的缺陷。第一，各自独立的评分使得传统的排序学习算法无法建模跨文档交互和捕捉局部上下文信息。已有的研究表明，融入局部上下文信息可以显著提高排序模型的有效性。第二，PRP逐个文档工作，而排序结果则应逐个查询进行评估。实际上，搜索引擎用户通常会在生成点击操作之前比较结果页面上的多个文档。此外，来自同一排序列表中其他文档的信息可能会影响标注者对当前文档的判断，这挑战了以下基本假设：相关性应该针对单个信息请求在每个文档上独立建模。
近年来，研究人员提出了一种新的排序学习算法，它使用了新的评分函数——多变量评分函数，以多个文档为输入，联合预测多个文档的排序分数。这样的方法可以捕获局部上下文信息，但对于输入文档的顺序十分敏感。这是因为他们都假设输入文档是一个序列，使用了DNN或RNN对其进行建模。当初始排序的效果不佳或意外地被打乱后，这些模型的效果也快速下降。
受到Set Transformer的启发，在本文中作者提出了一种新的多变量排序模型SetRank，它的输入是任意大小的文档集合，输出是文档的排列。SetRank的工作流程如下：它接受整个文档的集合作为输入，然后使用堆叠的多头自注意力块学习整个文档集合的嵌入表示，最后经过前馈神经网络输出文档的排序分数。
SetRank有着如下的优点：第一，与现有的多变量评分函数相似，它通过自注意力机制将整个文档集合作为输入，建模了文档间的关系；第二，SetRank学习了排列不变性函数作为其排序模型，因此输入文档的任意排列不会影响最终的排序结果。同时，由于自注意力是加性函数，它对输入文档集合的大小不敏感；最后，由于使用了序数嵌入向量，SetRank可以使用多个文档排序作为其初始排序。
Ranking Model Defined On Sets
给定查询 $q$ 和检索得到的文档集合 $D=\left[d_{i}\right]_{i=1}^{N} \subseteq \mathbb{D}$，排序的任务是为文档集合 $D$ 找到一个排列 $\pi \in \Pi_N$，使得某些效用最大化，其中 $\Pi_N$ 是索引 ${1,2,\cdots,N}$ 的所有排列的集合。
在排序学习中，对于每个查询 $q$，以此检索到的文档和对应的相关性标签记为 $\psi_{q}=\{D=\{d_{i}\}, \mathbf{y}=\{y_{i}\} \mid 1 \leq i \leq N\}$，$y_i$ 是文档 $d_i$ 的相关性标签。因此，整个数据集记作 $\Psi=\{\psi_q\}$，排序学习的目标是最小化如下的经验损失来产生最优的排序函数 $F(\cdot)$：
\[\mathcal{L}(F)=\frac{1}{|\Psi|} \sum_{\{D, \mathrm{y}\} \in \Psi} l(\mathrm{y}, F(D))\]其中 $\ell(\cdot)$ 是损失函数；$F(\cdot): \mathbb{D}^{N} \mapsto \mathbb{R}^{N}$ 是评分函数，它为每个文档计算得分，因此文档的排列结果可以通过对该评分进行降序排列得到：$\hat{\pi}=\operatorname{sort} \circ F(D)$
SetRank Model
如图所示，SetRank中的文档排序流程包括三步：表示，编码和排序。首先，表示层将每个输入的文档分别表示为特征向量，例如传统排序学习模型中使用的手工特征。同样，它可以通过序数嵌入向量将文档的初始排名包含在特征表示中。初始排名可以通过现有的排序模型（例如BM25或经过训练的LambdaMART模型）生成。其次，编码层通过包含相关文档的其他特征向量来丰富每个查询文档对的特征向量。在本文中，作者使用多头自注意力块（MSAB）或归纳多头自注意了力块（IMSAB）来获取一组查询文档对的表示形式作为其输入，并使用自注意力机制来生成一组新的表示形式。 MSAB或IMSAB块的多个子层以相同的结构堆叠在一起，用于对文档之间的高阶交互进行建模。第三，排序层接收最顶端的MSAB（或IMSAB）块的输出向量，将其传递给前馈（rFF）函数，生成所有文档的相关性得分，最后根据这些分数对文档进行排序。
文档表示
给定查询 $q$ 和关联的文档集合 $D=[d_1,d_2,\cdots,d_N]$，$D$ 中的每一个文档都可以表示为一个特征向量
\[\mathbf{d}_{i}=\phi\left(q, d_{i}\right), \text { where } \mathbf{d}_{i} \in \mathbb{R}^{E}\]其中 $\phi$ 是特征抽取的函数。它在传统的排序学习模型中是BM25或TF-IDF特征，本文中使用了基准数据集提供的特征。
除了传统的排序学习特征外，SetRank还可以利用文档的序数嵌入向量作为输入。在真实的搜索引擎中，关联文档可能具有由默认的排序模型（例如BM25或LambdaMART等）生成的一些先验的排序，因此作者提出使用序数嵌入函数 $P$ 将文档的绝对排序位置编码为与 $d_i$ 同等维度的向量：
\[\mathbf{p}_{i}=P\left(\operatorname{rank}\left(d_{i}\right)\right), \text { where } \mathbf{p}_{i} \in \mathbb{R}^{E}\]其中 $\operatorname{rank}(d_i)$ 代表 $d_i$ 由BM25或LambdaMART等模型生成的初始排序中的绝对位置。
之后文档的序数嵌入向量将分别与排序学习的特征相加得到 $N$ 篇检索文档的特征矩阵 $X \in \mathbb{R}^{N \times E}$：
\[\mathbf{X}=\left[\mathbf{d}_{1}+\mathbf{p}_{1}, \mathbf{d}_{2}+\mathbf{p}_{2}, \cdots, \mathbf{d}_{N}+\mathbf{p}_{N}\right]^{T}\]值得注意的是，某些时候可能不止一个初始的排序位置，因为可能会同时使用多个排序模型（例如BM25和LM4IR），因此多个不同的排序模型得到的序数嵌入向量可以直接相加得到总共的序数向量。SetRank允许输入多个初始排序。
文档编码
编码层的核心作用是将文档表示 $X^0 = X \in \mathbb{R}^{N \times E}$ 编码为内在的编码 $X^{N_b} \in \mathbb{R}^{N \times E}$。本文中使用了两种不同的编码方式，分别为MSAB和IMSAB。如下图所示。
MSAB基于DNN中的自注意力方法，它的形式是
\[\operatorname{Attn}(\mathrm{Q}, \mathrm{K}, \mathrm{V})=\operatorname{softmax}\left(\frac{\mathrm{QK}^{T}}{\sqrt{E}}\right) \mathrm{V}\]为了使得注意力函数更灵活，通常会使用多头策略，即
\[\text { Multihead }(\mathrm{Q}, \mathrm{K}, \mathrm{V})=\operatorname{concat}\left(\left[\operatorname{Attn}\left(\mathrm{Q} \mathbf{W}_{i}^{Q}, \mathrm{~K} \mathbf{W}_{i}^{K}, \mathrm{~V} \mathbf{W}_{i}^{V}\right)\right]_{i=1}^{h}\right)\]由此得到
\[\begin{aligned} \operatorname{MAB}(Q, K, V) &=\text { LayerNorm }(B+r F F(B)) \\ \quad \text { where } B &=\text { LayerNorm }(Q+\text { Multihead }(Q, K, V)) \end{aligned}\]其中 $rFF(\cdot)$ 是前馈神经网络层，$\text { LayerNorm }(\cdot)$ 是层归一化。最终的MSAB是
\[\operatorname{MSAB}(\mathbf{X})=\operatorname{MAB}(\mathbf{X}, \mathbf{X}, \mathbf{X})\]MSAB的缺陷是对于输入的集合大小较为敏感。在每个查询对应 $N_0$ 篇文档的数据集训练得到的MSAB在每个查询对应 $N_1$ 篇文档时表现不佳。为了解决这个问题，作者使用了IMSAB作为编码方式。
IMSAB构建 $M$ 个伪造的query向量 $I \in \mathbb{R}^{M \times E}$ 用于从原来的 $X \in \mathbb{R}^{N \times E}$ 中抽取特征 $H \in \mathbb{R}^{M \times E}$，它是 $M$ 个簇的中心。然后原始的文档 $X \in \mathbb{R}^{N \times E}$ 基于 $H$ 寻找上下文表示。
\[\operatorname{IMSAB}_{M}(X)=\operatorname{MAB}(X, H, H), \text { where } H=\operatorname{MAB}(I, X, X)\]文档排序
将多个MSAB（IMSAB）堆叠起来，将其结果通过前馈神经网络，就可以得到评分函数
\[\begin{array}{l} \mathbf{X}_{\mathrm{MSAB}}^{N_{b}}=\underbrace{\operatorname{MSAB}\left(\operatorname{MSAB} \ldots\left(\operatorname{MSAB}\left(\mathbf{X}^{0}\right)\right)\right)}_{N_{b}} \\ \operatorname{SetRank}_{\operatorname{MSAB}}(D)=\operatorname{rFF}_{1}\left(\mathbf{X}_{\mathrm{MSAB}}^{N_{b}}\right) \end{array}\]IMSAB的评分函数于此类似。最终的评分为
\[\hat{\pi}_{\mathrm{MSAB}}=\text { sort } \circ \text { SetRank }_{\mathrm{MSAB}}(D)\]模型训练
给定标注的查询 $\psi_{q}=\{D=\{d_{i}\}, \mathbf{y}=\{y_{i}\} \mid 1 \leq i \leq N\}$，其中 $y_i$ 是文档 $d_i$ 对于查询 $q$ 的信息增益，最优的注意力分配策略为
\[a_{i}^{y}=\frac{\tau\left(y_{i}\right)}{\sum_{d_{k} \in D} \tau\left(y_{k}\right)}, \quad \tau(x)=\left\{\begin{array}{cc} \exp (x) & x>0 \\ 0 & \text { otherwise } \end{array}\right.\]同理，给定预测的评分函数 $\{s_{1}, \cdots, s_{N}\}=\text{SetRank}(D)$，预测的注意力分配策略为
\[a_{i}^{s}=\exp \left(s_{i}\right) / \sum_{d_{k} \in D} \exp \left(s_{k}\right)\]因此交叉熵损失定义为
\[\mathcal{L}=\sum_{d_{i} \in D} a_{i}^{y} \log a_{i}^{s}+\left(1-a_{i}^{y}\right) \log \left(1-a_{i}^{s}\right)\]请注意，SetRank对输入集合的大小不敏感，尤其是对于基于IMSAB的SetRank。 在文档表示步骤中，可能需要使用序数嵌入向量来利用初始排名。 但是，由于无法使用较大的序数嵌入，它使模型无法处理大于训练数据中最大集合大小的输入集合大小。 为了解决这个问题，作者在模型训练期间提出了一种相对序数嵌入采样策略。首先，在训练时选择一个最大的集合大小 $N_{max}$，对于每个查询，对其起始数字进行采样得到 $s$，然后构建索引列表 $[s,s+1,s+2,\cdots,s+N-1]$ 代替原来的排序索引，这样所有小于 $N_{max}$ 的位置都可以训练。在测试时，大于 $N$ 小于 $N_{max}$ 的位置也是允许的。
