Jina AI将LLM Reranker延迟打下来了:21秒变3秒!
继Jina Reranker v2之后,Jina AI又开源了PE-Rank,一种新的基于LLM的重新排序器,用于高效的列表式段落重新排序。不是将原始文本输入到LLM的上下文窗口,而是使用嵌入模型将每个段落表示为一个特殊标记,然后将[指令]+[查询]+[特殊标记]输入到LLM中。在推理时,PE-Rank将输出空间限制在这些特殊标记上,从而实现更高效的解码。PE-Rank将重新排序100个文档的延迟从21秒大幅减少到3秒。
RankGPT(上)与PE-Rank(下)的比较。RankGPT将整个段落作为输入,并输出有序数字,而PE-Rank将特殊标记列表作为输入和输出。在右侧,展示了使用不同形式的输入在DL19上重新排序的结果。
PE-Rank背景
使用大型语言模型(LLM)作为重新排序器有什么吸引人的地方?
- 新任务的灵活指令
- 零样本能力
- 上下文推理
但在实践中,是什么阻止人们使用LLM作为重新排序器?
- 上下文长度:重新排序100个文档,每个文档有1000个标记,基本上需要一个10万标记的上下文长度。
- 大海捞针:由于重要信息可能在长上下文中丢失,性能可能会波动。
- 容易受到提示注入的影响:指令和查询可能会被候选文档覆盖。
- 输出格式问题:很难确保输出以正确的顺序格式,即d1 > d3 > d2 > d7。有时你会得到语法错误或过于冗长的结果。
PE-Rank原理
使用PE-Rank,LLM的输入基本上是指令+查询+嵌入的段落,每个都作为一个特殊标记,如下所示:
用于学习排名阶段训练的数据格式
其次,使用特殊段落标记来表示原始文本的想法类似于软提示。但在PE-Rank中,使用像Jina/BGE这样的外部嵌入模型来编码文档。这在外部嵌入和主干LLM自己的标记嵌入之间引入了一些差异,这意味着需要学习一个映射函数。为了做到这一点,冻结了嵌入模型和LLM,并仅训练了一个2层的多层感知器(MLP)来转换嵌入空间。
在两阶段排名范式下的 PE-Rank 概览。(a) 是检索阶段,检索 n 个段落嵌入;(b) 是大型语言模型(LLM)的前向传递过程;(c) 展示了列表式解码过程。
但是,如何微调一个大型语言模型(LLM)呢?经典的监督式微调(SFT)方法在这里有用吗?实际上,并不是特别有用。因为解码空间仅限于特殊的段落嵌入标记,应用标准的SFT并不直接。在PE-Rank中,结合了两种损失:ListMLE,它最大化生成下一个最相关段落标记的概率;以及上下文ListMLE,它进一步基于原始内容本身进行条件化。这种方法提高了模型利用查询和段落之间标记级交互的能力,并有助于在仅使用嵌入进行排名时转移这种能力。
两种类型的训练数据和学习排名的训练过程的说明
PE-Rank效果评测
使用 Mistral-7B-Instruct-v0.2 作为 PE-Rank 的大型语言模型(LLM)的基础模型,以及 Jina-embeddings-v2/BGE-v1.5 作为外部嵌入。通过这种设置,实现了与将原始文档投入 GPT-4(表中的 RankGPT4)相当的性能,但延迟仅为其六分之一,将总时间成本从20秒减少到3秒。如果只重新排序前20个候选项,单个查询的延迟可以进一步降低到0.5秒,使其在现实世界的应用中变得相当实用。
在TREC DL和BEIR上重新排序前100个段落的结果(NDCG@10)。Ret表示第一阶段使用的检索模型。
在推理过程中重新排序前100名候选者在不同阶段的延迟
在 Jina 嵌入和 BGE 嵌入之间切换时,发现 PE-Rank 可以一致地改进基础检索器的性能,无论是 BM25、Jina 还是 BGE。一个有趣的点是,尽管 BGE 在 MTEB 上的得分高于 Jina,但使用 BGE 嵌入重新排序 BM25 检索结果的性能在三个不同的数据集上始终低于使用 Jina 嵌入。这表明在像 MTEB 这样的通用嵌入基准测试中表现出色的模型,在这个特定上下文中可能不一定表现良好,而 Jina 嵌入在这里显示出更好的扩展性。
https://github.com/liuqi6777/pe_rank
https://arxiv.org/pdf/2406.14848
Leveraging Passage Embeddings for Efficient Listwise Reranking with Large Language Mo
本文转载自PaperAgent