FIRST: Faster Improved Listwise Reranking with Single Token Decoding

论文题目： FIRST: Faster Improved Listwise Reranking with Single Token Decoding

论文作者： Revanth Gangi Reddy

大语言模型也可以应用于信息检索（搜索）领域，比如精排，针对Listwise排序方法，将用户搜索query和候选passages输入大模型可以让大模型直接输出排序结果，一般来说，在每个passage前面可以给定一个标识符，比如A，B，C，D等，大模型可以输出排序后的passages标识符，比如BDCA表示排序顺序为B>D>C>A。大模型由于具有广泛的通用知识和泛化能力，因此在少样本甚至零样本的条件下也可以具有不错的效果，在特定领域可以进行微调以获得更好的效果。

但是，这种生成式的训练方法往往忽视了位置的相对信息，rerank往往更加关注排在前面的passages的准确性，因此排在前面出错的损失应该比排在后面排序错误的损失更大，即如果靠在前面的passage出现判断错误，代价应该更高。但是传统的训练方法采用CrosEntropyLoss损失函数无法建模这种损失关系。文本提出的FIRST方法，可以直接使用大模型输出的第一个logit进行候选passages的排序，不必输出所有的候选passage标识符，可以加快50%的推理速度并且可以取得一个不错的效果。此外，本文采用learning-to-rank损失来使得大模型关注更相关的passages的排序准确率。

论文的创新之处在于：

（1）介绍了FIRST方法，只通过大模型输出的第一个token的logit就可以完成候选passages的排序，可以大大加快推理速度。

（2）模型训练引入了learning-to-rank loss 提升了模型的效果。

（3）表明了LLM reranker在对检索召回进行相关性反馈（relevance feedback）方面具有应用潜力，对比cross-encoder排序模型能够取得更好的蒸馏效果。

一、背景介绍

搜索有两个非常重要的步骤，召回和精排，召回需要根据query召回多个候选的passages，比如通过双塔模型，可离线将所有的passages进行向量计算并进存储在向量数据库，在线只计算query的向量，然后检索出一批和query尽可能相近的passages，此过程往往也可以结合item召回等，提高召回的全面性。在精排阶段，就需要精确地判断和query最为相似的topk passages，此过程通过采用单塔模型（cross-encoder model），但是这种模型往往需要大量的标注数据进行监督训练才能获得一个不错的效果，很多工作已经结合大模型来优化精排模型。本文提出的FIRST实际上就是一种基于大模型实现精排的方法。

1.1 Learning to Rank（LTR）

在信息检索领域，Learning to Rank (LTR)目标是将候选item按照和query的相关程度进行排序，LTR是一类技术方法，LTR并不太关注每个item具体得的得分，更关注所有候选item的相对顺序。LTR最常见的应用是搜索场景，LTR有多种实现方法，可以大致分为三种主要方法：、pointwise（单点法）、pairwise（配对法）、Listwise（列表法）。

（1）pointwise排序。pointwise排序只同时考虑query和一个候选item之间的相关得分，不考虑item之间的关系，然后基于每个item计算出来的得分进行排序，可以得到候选item整体的顺序，也可以通过分类任务进行相关性的判断，这种训练任务相对简单。一组训练数据形式为<query, item_i, label_i>，i=1,2...N，N为候选item数量。

（2）pairwise排序。pairwise排序将排序任务转换成多个二元分类任务，只考虑两个item之间的关系。输入是query和两个item，输出是两个item的相对顺序。一组训练数据形式为<query, item_i, item_j, label_ij>，label_ij可以设置为1和-1，分别表示item_i>item_j和item_i<item_j。

（3）Listwise排序。pointwise和pairwise排序都没有考虑item处在搜索列表中的位置，Listwise排序需要根据query将所有的候选passages作为一个整体进行排序，排序需要考虑所有的候选passages进行一次性排序。它是将每个query和它所对应的所有搜索结果列表作为一个训练实例，训练复杂度相对较高。根据训练样例训练得到最优评分函数F，对应新的查询，评分F对每个item打分，然后根据得分由高到低排序，即为最终的排序结果，它可以考虑整体序列，针对特定的rank评价指标（比如MAP、NDCG）进行优化。一些相关工作也对此进行了优化工作。

1.2 基于大模型的Listwise排序

大模型具备较为强大的理解和推理能力，在预训练阶段学习了大量的世界知识，具有强大的泛化能力，这些知识通常是小模型所不具备的，因此可以通过大模型进行SFT让模型学习特定领域的知识才达到一个更好的效果。许多工作基于专有大模型，如GPT3.5/GPT4，进行Listwise排序，通过构造合适的prompt直接将query和候选的passages输入大模型，让大模型输出排序后的passages的标识符，利用大模型将Listwise排序当作一个生成任务，但是这种生成方式比较简单没有考虑候选passages之间的相对顺序，基于大模型实现Listwise也会导致prompt较长，因为需要一次性将所有的候选item都输入到模型中去，同时大模型由于参数量较大也会使推理速度变慢，在很多业务场景中往往不能承接耗时需求，这也是大模型在搜索领域进行落地的痛点之一。目前，可以通过大模型蒸馏小模型实现大模型在搜索领域的应用，提升精排模型的效果。

给定一个候选passages列表 $\mathcal{P} = \{p_{1},p_{1},...,p_{n}\}$ ，rerank的目标就是找出和query最相关的 $k$ 个passages，基于大模型的Listwise排序常常会采用滑动窗口策略来解决输入长度受限的问题，一个滑动窗口内有 $m$ （ $m<n$ ）个passages，设置步长（step size）为 $s$ ，每个滑动窗口内的passages都有一个唯一的标识符 $t_{i}$ ，LLM reranker生成按照相关程度降序排序的passages的标识符，比如 $t_{2} > t_{1} > t_{4} > t_{3}$ 。如下图所示，需要排序8个候选passages，滑动窗口设为4，步长设为2，具体过程是从后往前滑动进行窗口内相关性排序，比如蓝色部分表示前两个窗口，黄色部分代表最后一个窗口，步长为2表示上一个窗口的top2 passages将会参加下一个窗口内的排序。 $^{[1]}$

图片来源：arxiv.org/pdf/2304.09…

一些工作已经针对开源的LLM进行微调得到Listwise LLM reranker，监督微调方式是让模型按照和query的相关程度生成候选passages的标识符序列 $y = [y_{i}] > [y_{j}] > ... >[y_{m}]$ ， $y_{i}$ 和 $y_{j}$ 都是passage的标识符，其中对于query而言， $y_{i}$ 比 ${y_{j}}$ 更相关。使用语言建模目标对reranker进行训练，以最小化预测生成序列中下一个真实token的误差：

\ell_{LLM} = \sum_{i=1}^{|y|}log(P_{\theta}(y_{i}|x,y_{<i}))

二、FIRST

如上图所示，不同于之前的方法，FIRST期望通过LLM输出的第一个token的logits来对所有的候选passages进行排序，按照候选passages的唯一标识符（identifier）对应的logit值从大到小进行排序，从而得到整个passages列表的相对顺序。因此，推理时也只需计算一个token的logits就可以得到排序结果，大大节省了推理时间。因为排序是基于LLM词表中的单个token的logit，因此有必要避免将候选passage的标识符划分为多个token（tokenizer过程），本文中没有使用数字标识符，因为当>10时，基于BPE的tokenizer可能会将标识符划分为多个token。本文采用的大写字母作为passage的标识符，如A-Z，由于使用了滑动窗口机制，因此单个输入最大的候选passages数量不会超过20。

为了更好的基于第一个token的logits建模多个候选passages之间的相对顺序，需要针对性的提出合适的损失函数来优化模型，使用语言模型建模的交叉熵损失是次优的，因此本文提出了the weighted RankNet loss。

给定m个候选passages $(p_{1},p_{2},...,p_{m})$ ，t_{i}表示 $p_{i}$ 的标识符， $s_{i}$ 表示输出的第一个logits向量中 $t_{i}$ 对应的token所对应的logit值， $r_{i} \in [1,2,...,m]$ 是 $p_{i}$ 真实的排序顺序。the weighted RankNet loss表示为

\ell_{Rank}=\sum_{i=1}^{m}\sum_{j=1}^{m}\frac{I_{r_{i}<r_{j}}}{i+j}log(1+exp(s_{i}-s_{j}))=\sum_{r_{i}<r_{j}}\frac{1}{i+j}log(1+exp(s_{i}-s_{j}))

$\frac{1}{i+j}$ 为候选对 $(i,j)$ 的逆平均顺序（inverse mean rank），排在前面的权重更大，因此越重要。本文也结合了语言模型建模的交叉熵损失，最终的训练损失joint loss为

\ell_{Joint} = \ell_{LLM} + \lambda\ell_{Rank}

$\lambda$ 为超参数，控制两个损失的相对重要程度。训练过程中， $\ell_{Rank}$ 只作用于输出的第一个token的logits，而 $\ell_{LLM}$ 作用于所有的标识符token排序序列，在推理过程中则只基于输出的第一个token的logits获得所有候选passages的排序顺序。

三、实验

3.1 实验设置

（1）模型和参数设置：选择Zephyr $_{\beta}$ 作为初始模型，训练3个epoch，batch size设为32，学习率为5e-6，使用bf16精度训练，并且加入了noisy embeddings $^{[2]}$ ，使用deepspeed框架，在4张40G的A100卡上训练了大概7个小时，在MS Marco中随机采样了300个query作为验证集， $\lambda$ 设置为10。

（2）数据集：使用GPT-4标注的40k数据，共包含5k个query来自MS MARCO $^{[3]}$ ，样本包含不等的候选passages数量，但是都不超过20个。评估的时候使用BEIR benchmark $^{[4]}$ ，包含了MS MARCO中的数据以及其他领域的一些数据，比如医学、金融、维基百科等。

（3）rerank设置。使用Contriever $^{[5]}$ 检索候选passages，检索出来的top100 passages会作为reranker的输入，使用滑动窗口机制，窗口大小设置为20，步长设为10。

（4）对比基线（baseline）。对比了cross-encoder rerank $^{[6]}$ 、RankVicuna $^{[7]}$ 以及RankZephyr $^{[8]}$ 。后两个为基于大模型的listwise reranker，cross-encoder rerank使用500k从MS MARCO人工标注的pairwise数据，RankVicuna使用RankGPT数据进行训练 $^{[1]}$ ，包含了来自MS MARCO 100k的query数据并使用GPT-3.5进行标注。RankZephyr采用了两阶段训练，第一阶段使用RankGPT的数据，第二阶段使用了GPT-4标注的5k数据（query仍然来自MS MARCO）。

3.2 实验结果

（1）与基线的对比结果表1（table 1）为与基线模型的对比结果，从实验指标可以看出，虽然RankZephyr使用了更多的训练样本量（10w+5k），FIRST在大多数数据集上都超过了RankZephyr。cross-encoder在MS MARCO测试集上表现出了最佳效果，主要是因为它使用了人工在MS MARCO上标注的50w高质量数据。

（2）消融实验表2（table 2）为消融实验结果，主要结论如下：

使用joint loss相比只使用交叉熵loss效果更好
在RankNet loss的基础上添加逆平均顺序是有效的
只使用weighted RankNet loss要比只使用交叉熵loss表现要差，猜测可能是因为和大模型预训练目标不适配

（3）和其他rank loss的对比

此外，作者将RankNet loss和LambdaRank $^{[9]}$ loss以及ListNet $^{[10]}$ loss进行了对比实验，实验表明RankNet loss效果更好。

3.3 性能对比

本文提出的FIRST由于预测的时候只需要根据第一个token的logits就可以得到排序结果，不需要生成整个passages标识符序列，因此在推理延迟性方面，理论上应该具有更好的性能。图4在FIRST和生成整个标识符序列（Generation）方法上进行了推理时间上的对比，从图中可以看出，随着窗口的增大，FIRST的推理性能越明显，在窗口为20的时候推理速度差不多提升了2倍，因为Generation方法随着窗口的增大需要输出的token序列长度也会增加，从而导致时间的增加，而FIRST只基于大模型输出的第一个token的logits进行结果的计算。

图3为固定候选集数量为 $(20,30,40,60,80)$ ，然后计算FIRST在相应候选集数量上的延迟，之后在相同延迟下计算Generation方式所计算得到的passage数量，一共在6个数据集上进行了试验，结果证明了FIRST对比Generation方式在推理性能上的高效提升。

3.3 基于大模型reranker的相关性反馈

前面通过实验证明了基于LLM的reranke在排序性能上比cross-encoder更优，这对下游应用也是有帮助的。具体来说，可以考虑提供相关反馈（Relevance Feedback），以提高检索召回效果。在推理过程中使用reranker进行相关反馈涉及在测试时使用reranker对检索结果的输出优化检索器的query表示。Reddy等人 $^{[11]}$ 和Sung等人通过梯度下降基于query向量和cross-encoder reranker评分分布之间的KL散度损失来更新稠密检索器（如Contriever）的查询表示。由于reranker通常比检索器更强，使用更新后的query表示进行第二阶段检索时，可以提高之前检索结果的召回性能。

尽管cross-encoder reranker提供的得分可以用作蒸馏监督，listwise reranker输出的是候选项的有序序列。因此，在这种情况下，通常使用的KL散度损失无法用于相关反馈。对此，本文研究了如何利用listwise reranker进行相关反馈，以及它们是否能够相比cross-encoder在第二阶段检索召回效果上提供更大的改进。本文通过使用加权的RankNet loss进行实验，将listwise reranker的排名次序用作相关反馈的蒸馏监督。

实验结果表明来自LLM reranker的相关反馈相比来自cross-encoder reranker的相关反馈要具有更好的效果，此外还发现将LLM reranker和cross-encoder reranker的相关反馈结合起来（CE+LLM）会效果更好。

四、结论

本文介绍了FIRST，一种新的基于LLM做Listwise排序的方法。FIRST利用第一个LLM输出的第一个token的logits 来对候选passages进行排序，而不是采用典型的生成方法生成整个候选passages的标识符序列。本文通过实验展示了这种方法在相同时间内进行排序的候选passage数量显著大于使用生成整个序列的方法，从而在延迟受限的情况下进行rerank时获得更大的提升。FIRST还在训练时引入了learning-to-rank loss进一步提升了模型的效果，会让模型优先考虑排在前面的passages的正确性。

五、参考文献

[1] Sun, Weiwei, et al. "Is ChatGPT good at search? investigating large language models as re-ranking agents." arXiv preprint arXiv:2304.09542 (2023).
[2] Jain, Neel, et al. "Neftune: Noisy embeddings improve instruction finetuning." arXiv preprint arXiv:2310.05914 (2023).
[3] Bajaj, Payal, et al. "Ms marco: A human generated machine reading comprehension dataset." arXiv preprint arXiv:1611.09268 (2016).
[4] Thakur, Nandan, et al. "Beir: A heterogenous benchmark for zero-shot evaluation of information retrieval models." arXiv preprint arXiv:2104.08663 (2021).
[5] Izacard, Gautier, et al. "Unsupervised dense information retrieval with contrastive learning." arXiv preprint arXiv:2112.09118 (2021).
[6] Nandan Thakur, Nils Reimers, Johannes Daxenberger, and Iryna Gurevych. 2021a. Augmented SBERT: Data augmentation method for improving bi-encoders for pairwise sentence scoring tasks. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 296–310, Online. Association for Computational Linguistics.
[7] Ronak Pradeep, Sahel Sharifymoghaddam, and Jimmy Lin. 2023a. Rankvicuna: Zero-shot listwise document reranking with open-source large language models. arXiv preprint arXiv:2309.15088.
[8] Ronak Pradeep, Sahel Sharifymoghaddam, and Jimmy Lin. 2023b. Rankzephyr: Effective and robust zeroshot listwise reranking is a breeze! arXiv preprint arXiv:2312.02724.
[9] Christopher Burges, Robert Ragno, and Quoc Le. 2006. Learning to rank with nonsmooth cost functions. In Advances in Neural Information Processing Systems, volume 19. MIT Press.
[10] Zhe Cao, Tao Qin, Tie-Yan Liu, Ming-Feng Tsai, and Hang Li. 2007. Learning to rank: from pairwise approach to listwise approach. In Proceedings of the 24th international conference on Machine learning, ICML ’07, pages 129–136, New York, NY, USA. ACM.
[11] Revanth Gangi Reddy, Pradeep Dasigi, Md Arafat Sultan, Arman Cohan, Avirup Sil, Heng Ji, and Hannaneh Hajishirzi. 2023. Inference-time re-ranker relevance feedback for neural information retrieval. arXiv preprint arXiv:2305.11744.