【论文下饭】List-aware Reranking-Truncation Joint Model for Search and Retrieval-augmented Generation

List-aware Reranking-Truncation Joint Model for Search and Retrieval-augmented Generation

2024(GenRT)List-aware Reranking-Truncation Joint Model for Search and Retrieval-augmented Generation
期刊会议：WWW 2024
github: https://github.com/xsc1234/GenRT/

1 Introduction

列表感知检索通过捕获列表级的上下文特征以返回更好的列表，主要包括 重排序 和 截断。
• 重排序 会对列表中的文档进行精细重新评分。
• 截断 动态地确定排名列表的截止点，用于提高整体相关性和避免无关文档的错误信息。

图1 重排序 与 截断 拆分做的问题。在重排序中结果相似，但在截断中相差很大。拆分的流程导致误差累计问题 与 相关文档的丢失。

存在的问题：
问题一：重排序 和 截断 应该是相互依赖的任务，拆分的流程使得难以利用二者共享的信息。

• 重排序 建模的文档相关性，可以为 截断 提供重要基础。
• 截断 建模的文档相关性和位置之间的权衡，可以提供为 重排序 重要的上下文交互信息。

问题二：拆分的流程通常存在误差累计问题，重排序 的误差会影响到 截断，这在训练期间无法直接优化。

• 两个阶段的文档相关性判断标准不一致。
• 两个阶段的关注点不同。
  • 截断 对排在前面的文档较为敏感。因为如果有太多不相关的文档排在列表前面的话，会导致列表在这些文档处截断，从而丢失了排在它后面的相关文档。
  • 重排序 更加关注整个列表的排名情况，对不相关的文档排在列表的前面的情况并不敏感。
  • 图1中(a)和(b)有相似的重排序性能表现（0.89和0.90）。在(b)中有两个不相关的文档排在前面（B和E），这就导致了(b)的截断性能差于(a)（19.4<25.8）。

本文需要解决的问题：

• C1：如何共享 重排序 和 截断 任务的建模信息？
• C2：重排序 是一个动态改变排名列表的过程，但 截断 的决策需要基于一个静态列表，如何同时执行？
• C3：如何设计联合学习的损失函数？

图2 GenRT与之前方法的区别

如图2所示，本文提出GenRT，它是一个通过序列生成实现的重排序-截断联合模型。

C1：如何共享 重排序 和 截断 任务的建模信息？
GenRT设计了一个全局依赖编码器（global dependency encoder），提供 全局的列表级别上下文特征，以在重排序和截断之中共享。

C2：重排序 是一个动态改变排名列表的过程，但 截断 的决策需要基于一个静态列表，如何同时执行？
GenRT采用了序列生成范式，逐步输出最终的重排序列表。在每一步，根据前一个状态选择处于当前排名的文档，并同时做出局部最优的截断决策。截断被转换为一个分类任务，它根据列表的（排名）前后双向信息进行解决。序列生成范式记录了前向信息，我们还引入了局部后向窗口来提供后向信息。这样，我们的模型可以将动态重排序与静态截断结合起来。

C3：如何设计联合学习的损失函数？
GenRT设计了步骤自适应注意力损失（step-adaptive attention loss）和逐步lambda损失（step-by-step lambda loss），将它们作为重排序的目标函数。GenRT还引入了奖励增强最大似然（reward augmented maximum likelihood，RAML），设计了每一步截断的 基于RAML的软损失。

2 Related Work

2.1 Reranking in List-aware Retrieval

在列表感知检索中的重新排序利用列表级别的上下文特征来重新评分和排名列表中的每个文档。

• DLCM[1]
• GSF[2]
• SetRank[28]
• PRM[29]
• IRGPR[25]
• DASALC[32]
• SRGA[30]
• MIR[42]
• Seq2Slate[5]
• Globalrerank[48]

2.2 Truncation in List-aware Retrieval

截断的目标是确定输入排名列表的最佳截止点，以实现整体相关性和排除无关文档之间的最佳权衡。

• BiCut[23]
• Choppy[4]
• AttnCut[40]
• LeCut[26]

3 Method

图3 GenRT的总览。全局依赖编码器捕捉初始序列的特征。序列依赖解码器一步步生成最后的序列，该序列的文档相关性降序排列，并且同时实现截断决策。 图4 GenRT编码器和解码器架构。(a)全局依赖编码器；(b)序列依赖解码器（第T步）

3.1 Global Dependency Encoder

文档特征表示
全局依赖编码器捕获 列表级别的上下文特征。

文档编码成embedding后，作为 编码器的输入。
给定一个查询$q$，文档列表$D=[d_1,d_2,...,d_N]$ 和 初始排名分数列表$L=[l_1,l_2,...,l_N]$，文档$d_i$的embedding表示为：

$$u_{d_i}=f(q,d_i,l_i),u_{d_i}\in {\mathbb{R}}^Z,i\in [1,N]\text{\hspace{1em}(1)}$$

其中，$N$是输入列表中文档的数量，$f$用于融合$q$，$d_i$和$l_i$特征。

特别地，
对于基于特征排名任务（feature-based ranking tasks），例如MSLR（输入数据是LTR特征）。作者跟随SetRank，采用了传统的LTR方法来提取特征（matching，pagerank等），然后拼接上初始排名分数$l_i$得到${\mathbf{u}}_{d_i}$。

对于基于文本的排序任务，例如Natural Questions（输入数据是文本）。作者使用基于交互的排序模型（interaction-based ranking model）输出的[CLS]词元作为查询$q$和文档$d_i$特征表示，然后将初始排名分数映射到 可学习的位置编码（learnable position embedding）${\mathbf{lp}}_i\in {\mathbb{R}}^Z$上，并将其按元素相加 得到${\mathbf{u}}_{d_i}$。

这样，我们就得到了整个文档特征矩阵$\mathbf{U}\in {\mathbb{R}}^{N\times Z}$

$$\mathbf{U}=[{\mathbf{u}}_{d_1},{\mathbf{u}}_{d_2},...,{\mathbf{u}}_{d_N}{]}^T\text{\hspace{1em}(2)}$$

Transfer Layer
传递层（Transfer Layer）的作用是将$U$映射到特定的维度（对齐维度）。
$$\mathbf{X}=\text{Swish}(\text{MLP}(\mathbf{U})),\mathbf{X}\in {\mathbb{R}}^{N\times E}\text{\hspace{1em}(3)}$$

笔者注：
$$\text{Swish}(x)=x\cdot \text{Sigmoid}(x)$$

其中，$\text{MLP}$是多层感知机，$\text{Swish}$为激活函数。

Multi-head Self Attention

$$\mathbf{O}=\mathbf{X}+\text{MHSA}(\text{LN}(\mathbf{X}))\text{\hspace{1em}(4)}$$

3.2 Sequential Dependency Decoder

序列依赖解码器遵循序列生成范式，逐步生成按相关性从高到低的序列，并在每一步中基于双向的序列信息进行截断决策。在图4中展示了第T步的解码器操作。

交叉排序FFN（Cross Ranking FFN）和 动态排序模块（Dynamic Ranking Module）确定动态排名列表，并在每个步中选择最佳输出文档。同时，截断模块（Truncation Module）根据动态排名模块和 局部后向窗口（Local Backward Window）获得序列的双向信息 进行截断决策。

给定输入文档列表 D ′ = { r 1 1 , r 1 2 , . . . , r 1 T − 1 } D' = \{r_1^1, r_1^2, ..., r_1^{T-1}\} D′={r11​,r12​,...,r1T−1​}（$r_1^t$为第$t$步输出的文档），注意这些文档都已经排好序了。

Transfer Layer
同3.1，通过文档特征矩阵$U^{\prime }$，得到$X^{\prime }$

Multi-head Self-Attention
$$\text{MHSA}(X^{\prime })\to [{\mathbf{m}}_{r_1^1},{\mathbf{m}}_{r_1^2},...,{\mathbf{m}}_{r_1^{T-1}}]$$

Cross Ranking FFN
交叉排序FFN预测每一个文档的分数，从而在当前步中选择最佳的文档。它的输入来自解码器和编码器。

在解码器这边，${\mathbf{m}}_{r_1^{T-1}}$是第T步的序列依赖信息。扩张${\mathbf{m}}_{r_1^{T-1}}\in {\mathbb{R}}^E$为矩阵$\mathbf{M}\in {\mathbb{R}}^{N\times E}$，其每一行都是${\mathbf{m}}_{r_1^{T-1}}$。

在编码器这边，$\mathbf{I}$通过特征交叉（latent cross）[6][32]得到，其中$\text{FFN-Swish}$为MLP和Swish层，$\odot$为按元素乘操作。
$$\mathbf{I}=(1+\text{MLP}(\mathbf{O}))\odot \text{FFN-Swish}(\mathbf{U})\text{\hspace{1em}(5)}$$

[6] Latent Cross: Making Use of Context in Recurrent RecommenderSystems.
[32] Are Neural Rankers still Outperformed by Gradient Boosted Decision Trees?

$\mathbf{I}$为当前步的候选文档集特征矩阵，$\mathbf{M}$为当前步的序列依赖矩阵。$\mathbf{M}$与每个文档的特征$\mathbf{I}$进行交互，从而得到每个文档的分数$\mathbf{S}$。

$$\mathbf{S}=\text{rFFN}(\text{Concat}(\mathbf{I},\mathbf{M})),\mathbf{S}\in {\mathbb{R}}^N\text{\hspace{1em}(6)}$$

Dynamic ranking module
动态排名模块 遮蔽 步数在1到T-1的文档（避免选择重复文档），并根据降序的$\mathbf{S}$对剩余的候选文档进行排序，以获得当前步的排名列表 R = { r 1 T , r 2 T , . . . , r N − T + 1 T } R=\{r_1^T, r_2^T, ..., r_{N-T+1}^T\} R={r1T​,r2T​,...,rN−T+1T​}。

动态排名模块 选择排名第一的文档（$r_1^T$） 进行以下的流程。

Truncation Module
截断模块 将截断问题转换为（在每一步的）二分类问题。截断模块 通过聚合 排在$r_1^T$前面和后面的（文档）信息，来做出截断决策。

截断模块记录了1到T-1步的文档特征$\mathbf{O}$，记$\mathbf{F}=[{\mathbf{o}}_{r_1^1},{\mathbf{o}}_{r_1^2},...,{\mathbf{o}}_{r_1^{T-1}}{]}^T$。

重排序列是一步步生成的，当输出当前步的文档$r_1^T$时，模型无法捕获到排在$r_1^T$后面的文档的信息。为了解决这个问题，局部后向窗口（local backward window）取$\beta$个排在$r_1^T$后面的文档特征$\mathbf{B}=[{\mathbf{o}}_{r_2^T},{\mathbf{o}}_{r_3}^T,...,{\mathbf{o}}_{\beta +1}^T{]}^T$。

为什么选择$\beta$个文档，是因为当前排序列表是局部的排序，并不能代表全局，如果选择了所有文档，那么会带来噪声。

截断模块的输入为$\mathbf{G}=\text{Concat}(\mathbf{F},{\mathbf{o}}_{r_1}^T,\mathbf{B})$。

为了区分前向和后向的信息 以及 文档在$\mathbf{G}$中的位置，作者引入了类似T5[33]的相对位置编码。相对位置编码的多头自注意力为${\text{MHSA}}_{pos}$聚合序列双向的信息，并作出截断决策。

$${\text{MHSA}}_{pos}(\mathbf{G})\to [{\mathbf{j}}_{r_1^1},...,{\mathbf{j}}_{r_1^{T-1}},{\mathbf{j}}_{r_1^T},{\mathbf{j}}_{r_2^T},...,{\mathbf{j}}_{r_{\beta +1}^T}{]}^T$$

$$P=\text{Softmax}(\text{MLP}({\mathbf{j}}_{r_1^T})),\mathbf{P}\in {\mathbb{R}}^2\text{\hspace{1em}(7)}$$

其中，$\mathbf{P}=[p_0,p_1]$是一个二元概率分布，代表在当前步下，是否要进行截断。

如果截断，那么GenRT会直接返回1到T步生成的文档，否则整个流程将继续直到做出截断决策 或 排序完整个文档列表。

3.3 Training and Inference

作者针对 重排序 和 截断设计了不同的损失。

step-adaptive attention loss (sa-att loss)
设第T步中，给定一个查询$q$，待排序的文档列表$D=[d_1,d_2,...,d_N]$，其对应的相关性标签 Y = { y 1 , y 2 , . . . , y N } Y=\{y_1, y_2, ..., y_N\} Y={y1​,y2​,...,yN​}，以及已经选择的文档列表$D^{\prime }$。那么，真实标签的注意力分数$a_i$为
$$a_i=\frac{\exp (\varphi (d_i))}{{\sum }_{d_j\in D}\exp (\varphi (d_j))},\varphi (d_i)=\{\begin{array}{ll}-10^4,& d_i\in D^{\prime };\\ y_i,& \text{otherwise.}\end{array}\text{\hspace{1em}(8)}$$

笔者注：
原损失函数来自[1]，作者对此损失函数的改造主要在，对于已经选择的文档，赋予较小的注意力权重。

对于预测的标签分数$\mathbf{S}=[s_1,s_2,...,s_N{]}^T$，采用相同的注意力分数计算公式（同上），记为$b_i$。

第T步的损失为：

$$L_{sa-att}^T=-\sum _{d_i\in D}{a}_i\log (b_i),{\mathcal{L}}_{sa-att}=\sum _{t=1}^N{\alpha }_t{L}_{sa-att}^t\text{\hspace{1em}(9)}$$

其中，${\alpha }_t=\frac{1}{\log (1+t)}$为衰减权重，这是为了让模型更加关注排在前面的文档。

step-by-step lambda loss (sbs loss)
sbs loss对于那些 错排的文档对 进行惩罚。（例如，$y^{t_b}$大于$y^{t_f}$，那么$s_{t_b}$应该大于$s_{t_f}$，那么$r_1^{t_b}$应该排在$t_1^{t_f}$前面）。

$${\mathcal{L}}_{sbs}=\sum _{i=1}^{ϵ}\sum _{j=i+1}^{ϵ}\mathbb{I}(y^j>y^i)\Delta N\log (1+e^{s_i-s_j})$$

其中，$\Delta N$为Lambda Loss (${\text{NDCG}}_{\text{swap}}-\text{NDCG}$)。

(LambdaLoss Framework) The LambdaLoss Framework for Ranking Metric Optimization
期刊会议：ACM 2018
博客：
learning2rank: LambdaLoss Framework 学习

重排序的整体损失为$\mathcal{L}={\mathcal{L}}_{sa-att}+\eta {\mathcal{L}}_{sbs}$，其中$\eta$为超参数权重。

---

截断损失
作者基于RAML[27]设计了截断损失，基于[17]定义了截断评估指标TDCG。

[17] Cumulated gain-based evaluation of IR techniques
[27] Reward Augmented Maximum Likelihood for Neural Structured Prediction
博客：
有监督学习还是强化学习？小孩子才做选择，大人全都要

$$\text{TDCG}@x=\sum _{t=1}^x\frac{\gamma (y^t)}{\log (t+1)}\text{\hspace{1em}(10)}$$

笔者注：
论文中并没有说$\gamma$长什么样，但看了看代码，似乎不关键？

其中，$x$为截断点，$\gamma$为对低相关性文档的添加惩罚项（add penalty items）。高相关性的文档TDCG更高，反之TDCG则更低。这种奖励机制就要求模型尽可能的返回少的低相关，并返回尽可能多的相关文档。

给定已经排序好的文档列表 D ′ = { r 1 1 , r 1 2 , . . . , r 1 T − 1 } D' = \{r_1^1, r_1^2, ..., r_1^{T-1}\} D′={r11​,r12​,...,r1T−1​}，此时选择生成的文档为$r_1^T$，还剩下的文档（局部后向窗口）为 R ′ = { r 2 T , r 3 T , . . . , r β + 1 T } R' = \{r_2^T, r_3^T, ..., r_{\beta+1}^T\} R′={r2T​,r3T​,...,rβ+1T​}，局部排序列表为 { r 1 1 , r 1 2 , . . . , r 1 T − 1 , r 1 T , r 2 T , r 3 T , . . . , r β + 1 T } \{r_1^1, r_1^2, ..., r_1^{T-1}, r_1^T, r_2^T, r_3^T, ..., r_{\beta+1}^T\} {r11​,r12​,...,r1T−1​,r1T​,r2T​,r3T​,...,rβ+1T​}。
如果模型在当前步决定截断，此时奖励为$\text{TDCG}@T$，否则为$\text{TDCG}@(T+\beta )$，$\beta$为局部后向窗口的大小。在第T步的二元软标签为：

$$y_{cut}^T=\frac{\exp (\text{TDCG}@T)}{\exp (\text{TDCG}@T)+\exp (\text{TDCG}@(T+\beta ))}$$

$$y_{nocut}^T=\frac{\exp (\text{TDCG}@(T+\beta ))}{\exp (\text{TDCG}@T)+\exp (\text{TDCG}@(T+\beta ))}$$

截断损失如下：

$${\mathcal{L}}_T=-\sum _{t=1}^N(y_{cut}^t\log (p_1^t)+y_{nocut}^t\log (p_0^t))$$

其中，$p_1^t$和$p_0^t$在公式$(7)$中定义。

---

预热训练
在第一个epoch训练中，仅训练 重排序。让模型学习判断文档相关性的能力。在后续训练中，模型在每个批次（batch）中交替训练 重排序 和 截断。即，当训练重排序时，固定截断模块参数，仅使用${\mathcal{L}}_R$作为损失函数；当训练截断时，固定交叉排序FFN的参数，仅使用${\mathcal{L}}_T$作为损失函数。

GenRT还可以适用于仅需要 重排序 或 截断 的IR系统中。

对于仅需要 重排序 的情景，作者提出一种平衡延迟和准确性的加速策略。在推理截断，它直接使用可学习的start向量作为序列依赖，并使用在第一步的得分矩阵$\mathbf{S}$作为重排序的结果，而不需要生成序列。

4 Experiments

4.1 Experiment Settings

Datasets
(1) Learning-to-rank 排序学习数据集。这三个数据集都是从搜索引擎中得来的，并且都是feature-based。每个样本都是一个特征向量，标签有五个级别从0（不相关）到4（特别相关）。

• Microsoft LETOR 30K (MSLR30K)[31]
• Yahoo! LETOR set1 (Yahoo!)
• Istella LETOR (Istella)

(2) Open-domain Question-Answering 开放域问答数据集。这两个数据集都是text-based，标签有两个级别从0（不相关）到1（相关）。

• Natural Questions [22]
• TriviaQA [18]

Baslines and Evaluation Metrics
…
Implementation
…

4.2 Performance on Web Search

作者遵循SetRank[28]和DLCM[1]中的设置，使用RankLib实现的LambdaMart来检索每个查询的前40个文档作为输入排序列表。这些列表被用作列表感知重排模型的输入。

List-aware Reranking Performance

表1 不同 重排序 模型在learning-to-rank数据集上的表现（加粗：最好的结果；下划线：第二好的结果；十字：与baseline模型的结果显著性差异p值小于等于0.05）

笔者注：
在MLLR 30K、Istella上有较为显著的结果

List-aware Truncation Performance

表2 不同截断模型在learning-to-rank数据集上的表现

4.3 Performance on Retrieval-augmented LLMs

使用Wikipedia passage-collection作为语料，Contriever [16]和interaction-based ranker [8]作为检索-排序模块，获得前40个文章，送入排序模型。使用gpt-3.5-turbo-16k作为LLM。使用EM [34]来计算LLM的回答准确性。

List-aware Reranking Performance

表3 在检索增强的LLM上的重排序性能表现。

List-aware Truncation Performance

表4 在检索增强的LLM上的截断性能表现。Length为在列表中文章的数量。

笔者注：
重排一般般，截断还优点效果。

4.4 Analysis

图5 (a) 得分最小距离随步数T的变化。(b) TDCG随着局部后向窗口大小的变化（β）。

Effect of Generative Ranking
得分最小距离：在每一步的待选择文档中，最小的正样本得分 减去 最高的负样本得分。

GenRT的距离大于SetRank的距离，并且随着步骤𝑇的增加而增加，这表明序列依赖使得相关的文档更容易被选择。

Size of Local Backward Window
当$\beta$从0增加到4时，性能得到提升，因为引入了更多有价值的后向信息；当$\beta$大于4时，性能开始下降，因为引入了更多的噪声。

Loss Function in Reranking

表5 重排序 的两个损失函数的消融实验。

${\mathcal{L}}_{sa-att}$对重排序最为关键，${\mathcal{L}}_{sbs}$次之。

Efficiency Analysis of List-aware Retrieval
由于GenRT需要逐步生成列表，这就是导致了IR系统推理时间的增加。

为了缓解这个问题，作者提出了两种加速策略。

其中一种是，针对只需要重拍而不需要截断的场景，可以直接使用第一步（first step）得到的得分矩阵$\mathbf{S}$作为重排结果，称为${\text{GenRT}}_{\text{fast}}$

另一种是，使用比SetRank更少的参数量，特别是解码部分。作者对了在Tesla V100上模型在重排上的耗时。${\text{GenRT}}_{\text{fast}}$比SetRank快0.6倍，GenRT则慢2.1倍。对于重排-截断的耗时，GenRT比SetRank-AttnCut慢1.6倍。

模型	耗时倍率↓
SetRank	1x
${\text{GenRT}}_{\text{fast}}$	0.6x
GenRT	2.1x
------------------------------	----------
SetRank-AttnCut	1x
GenRT	1.6x

Appendix

A.3 Efficiency Analysis

表7 重排序-截断流程效率分析 表8 重排序-截断，然后进行检索增强LLM流程效率分析 表9 截断模型的性能表现，LTR数据集，SetRank作为reranker。