浅入浅出 Rerank 模型

引言
#

在构建基于 RAG 的大模型应用中，有一个绕不开的话题 Embedding、向量检索和 Rerank。

Embedding 和向量检索在之前的文章有所介绍。

向量相似性检索方法

2023-10-30·5606 字·12 分钟

检索算法 RAG 向量数据库算法原理

本文详细介绍了多种向量相似性检索方法，如 KD 树、IVF 倒排索引、HNSW 和 LSH。通过分析 Annoy、Faiss、PGVector 和 FALCONN 源码中的具体实现，从数据结构到算法实现进行了详细的介绍。

本文主要是想浅谈一下 Rerank 的原理，从 Transfomer 之后的发展历史，最后浅谈一下是否需要使用 Rerank。

Rerank 概念和应用场景
#

先说说 Rerank，直观地理解他就是对搜索结果的重排序。这个概念本身并不是什么新东西，Rerank 在传统的大型搜索引擎，甚至 APP 中的搜索中都有广泛地使用。例如：

百度搜索在检索内容之后会根据用户个性化进行 Rerank 重排序
大众点评会在服务端检索完成后，在客户端根据用户喜好进行 Rerank 重排序 (*端智能在大众点评搜索重排序的应用实践.[2022-02]*)

除此之外，在 NLP 自然语言处理领域也很早就有了 Rerank 这个概念，例如：

针对自然语言处理用于判别行的 Rerank 重排序 (*Discriminative Reranking for natural language parsing.[2005-03]*)

如果想对 Rank 有更多了解，可以看《Learning to Rank for Information Retrieval》

所以 Rerank 主要是用在大范围检索后的小范围精准排序，所以比起在大数据量下用倒排、向量或者词频等较为标准化的检索方式，Rerank 主要是针对特定的细节和关联关系做重排。比如上面提到的根据用户个性化重排等等。

回到 RAG 场景本身，向量检索（词嵌入 Embedding + 向量库）尤其在词嵌入 Embedding 的部分其实是对原始的自然语言是有损的，且不同的向量维度和不同的模型都会有不同程度的有损。至于全文检索（BM25）是使用分词+词频匹配的，本身也不具备语义检索。所以 Rerank 可以在检索完之后直接对原始自然语言进行更精准地排序。

但 Rerank 算法一般效率比较低，所以只适合小范围的精准排序，同时当前的许多 Rerank 算法也相当消耗计算资源。

基于 BERT 的 Rerank —— 编码器
#

从 17 年 Transformer 的提出，再到 18 年 Bert 和 GPT-1 模型。研究者们发现这种基于 Transformer 的预训练语言模型，在自然语言处理上取得了不错的进展。自然而然可以预见，对自然语言的 Rerank 也许可与走类似的道路。

2019 年 1 月有研究者提出了基于 Bert 的 Rerank 模型（*Passage Re-ranking with BERT.[2019.01]*）。论文中研究者将 Bert-LARGE 模型作为一个分类模型，即使用 BERT 中的 [CLS] 标记。通过单层神经网络输入从而获取概率输出，根据这些概率对文本进行重排序，类似于 Cross-Encoder。

$$ \mathbf{Input}:\space \mathrm{[CLS]} \space query\_token\_ids \space \mathrm{[SEP]}\space doc\_token\_ids\space \mathrm{[SEP]} $$

这部分在代码中也有直接的体现，主要是将数据集转换为 tensorflow Features 中的这段：

# https://github.com/nyu-dl/dl4marco-bert/blob/master/convert_msmarco_to_tfrecord.py

def write_to_tf_record(writer, tokenizer, query, docs, labels,
                       ids_file=None, query_id=None, doc_ids=None):
  query = tokenization.convert_to_unicode(query)
  # 转化 query token id，在前面添加了 [CLS]
  query_token_ids = tokenization.convert_to_bert_input(
      text=query, max_seq_length=FLAGS.max_query_length, tokenizer=tokenizer, 
      add_cls=True)

  query_token_ids_tf = tf.train.Feature(
      int64_list=tf.train.Int64List(value=query_token_ids))

  for i, (doc_text, label) in enumerate(zip(docs, labels)):
	# 转换为 doc token ids， 没有添加 [CLS]
    doc_token_id = tokenization.convert_to_bert_input(
          text=tokenization.convert_to_unicode(doc_text),
          max_seq_length=FLAGS.max_seq_length - len(query_token_ids),
          tokenizer=tokenizer,
          add_cls=False)

    doc_ids_tf = tf.train.Feature(
        int64_list=tf.train.Int64List(value=doc_token_id))
        
	# 数据集标签
    labels_tf = tf.train.Feature(
        int64_list=tf.train.Int64List(value=[label]))
	
    features = tf.train.Features(feature={
        'query_ids': query_token_ids_tf,
        'doc_ids': doc_ids_tf,
        'label': labels_tf,
    })

nyu-dl/dl4marco-bert

Python

480

最后需要提及的是，该论文团队尝试使用了 MS MARCO 数据集和 TREC-CAR 数据集对 Bert 进行了训练和评估，并取得了还不错的效果。

类似上面这种方式，另一篇论文 *Understanding the Behaviors of BERT in Ranking.[2019.04]*，又进行了进一步研究。提出了四种基于 BERT 的 Rank 方法，其中假设 $q$ 代表；$d$ 代表 document；$qd$ 代表 query 和 document 的拼接形式，两者之间会被标记 [SEP] 标志；$last$ 代表 Bert 中的最后一层即 $k = 24$：

BERT(Rep)：分别计算 query 和 document 的 Embedding 表示，然后计算两者 Embedding 的 cos 余弦距离。此方法类似于 Bi-Encoder。

$$ \mathbf{Input \space 1}:\space \mathrm{[CLA]} \space query\_token\_ids \space \mathrm{[SEP]} \\ \mathbf{Input \space 2}: \space \mathrm{[CLA]} \space doc\_token\_ids \space \mathrm{[SEP]} \\ \mathbf{BERT(Rep)}(q,d) = cos(\overrightarrow{q}_{cls}^{last}, \overrightarrow{d}_{cls}^{last}) $$

BERT(Last-Int)：和上面这篇论文一样，是 BERT 官方最推荐的，这里的 $w$ 指的是权重，即与 $w$ 的线性组合。此方法类似于 Cross-Encoder。

$$ \mathbf{Input}:\space \mathrm{[CLS]} \space query\_token\_ids \space \mathrm{[SEP]}\space doc\_token\_ids\space \mathrm{[SEP]} \\ \mathbf{BERT(Last\text{-}Int)}(q,d) = w^T \overrightarrow{qd}^{last}_{cls} $$

BERT(Mult-Int)：在方法上的基础上，添加不同 Bert 层的匹配特征，以此来研究 Bert 中不同层是否会提供不同信息。

$$ \mathbf{BERT(Mult\text{-}Int)}(q,d) = \sum_{1\le k \le 24} (w^k_{Mult})^T \overrightarrow{qd}^{k}_{cls} $$

BERT(Term-Trans)：算是 1+3 的增强版本，在 Bert 上添加一个神经排名网络。首先使用 query 和 document 之间的转换矩阵，计算它们的 Embedding 余弦距离；然后使用均值汇聚+线性组合来组合 Bert 所有层的转换矩阵。

$$ s^k(q,d)=\mathrm{Mean}_{i,j}(\cos(rule(P^k\overrightarrow{q}_i^k),rule(P^k\overrightarrow{d}_i^k)))\\ \mathbf{BERT(Term\text{-}Trans)}(q,d) = \sum_{1\le k \le 24} w^k_{trans}s^k(q,d) $$

最终结论表明基于 Bert 的 Rank 在 MS MARCO 的重排序中表现能力超过了目前的神经网络检索模型，即在问答场景中是很有效的。但是在 TREC Web Track 临时文档的排名能力较差，因为 TREC 风格文档的排序和场景中的用户点击有关系，而不是文本上下文本身。

还有一些在 Bert 场景下的经典论文这里不再复述了，如 [2019.10] Multi-Stage Document Ranking with BERT，感兴趣的可以看下。

基于 GPT 的 Rerank —— 解码器
#

Transformer 出现之后，基于 Transformer 架构来做 Rank 或 Rerank 的研究工作主要还是集中于以编码器为主的 Bert 模型。在 ChatGPT 出圈之后，部分研究者的视角也逐渐转移到以解码器为主的 GPT 模型，探讨如何使用解码器来进行 Rerank (*SGPT: GPT Sentence Embeddings for Semantic Search.[2022.02]*)。

Muennighoff/sgpt

SGPT: GPT Sentence Embeddings for Semantic Search

Jupyter Notebook

864

SGPT 提出了两种方式：

SGPT Cross-Encoder：字面意思就是通过 GPT 实现 Cross-Encode 交叉编码器。具体的方法是将 query 和 document 拼接起来一起编码，然后基于获取的对数概率 (log probabilities) 来计算分数。

## https://github.com/Muennighoff/sgpt/blob/main/README.md#asymmetric-semantic-search-ce

prompt = 'Documents are searched to find matches with the same content.\nThe document "{}" is a good search result for "'

for query in queries:
    print(f"Query: {query}")
    for doc in docs:
        context = prompt.format(doc)

        context_enc = tokenizer.encode(context, add_special_tokens=False)
        continuation_enc = tokenizer.encode(query, add_special_tokens=False)
        ## 拼接 query 和 document 
        model_input = torch.tensor(context_enc+continuation_enc[:-1])
        continuation_len = len(continuation_enc)
        input_len, = model_input.shape

        ## 获取对数概率
        logprobs = torch.nn.functional.log_softmax(model(model_input)[0], dim=-1).cpu()
        logprobs = logprobs[input_len-continuation_len:]
        ## Gather the log probabilities of the continuation tokens -> [continuation_len]
        logprobs = torch.gather(logprobs, 1, torch.tensor(continuation_enc).unsqueeze(-1)).squeeze(-1)
        score = torch.sum(logprobs)
        ## The higher (closer to 0), the more similar
        print(f"Document: {doc[:20] + '...'} Score: {score}")

SGPT Bi-Encoder：同理使用 GPT 来实现 Bi-Encoder 双编码器。具体的方法比较简单，即分别计算 query 和 document 的 Embedding，然后计算余弦距离。

## https://github.com/Muennighoff/sgpt/blob/main/README.md#asymmetric-semantic-search-be

## 分别计算 query 和 document 的 embedding 
query_embeddings = get_weightedmean_embedding(tokenize_with_specb(queries, is_query=True), model)
doc_embeddings = get_weightedmean_embedding(tokenize_with_specb(docs, is_query=False), model)

## 计算余弦距离，分数在 [-1, 1]
cosine_sim_0_1 = 1 - cosine(query_embeddings[0], doc_embeddings[0])
cosine_sim_0_2 = 1 - cosine(query_embeddings[0], doc_embeddings[1])
cosine_sim_0_3 = 1 - cosine(query_embeddings[0], doc_embeddings[2])

Rerank 模型盘点
#

最后让我们盘点一些知名且大规模使用的 Rerank 模型，如智源实验室的 BGE、阿里的 GTE、Jina AI、Cohere 等，看看他们是如何实现的。

需要注意的一点是，在上面我们已经可以看出来，Rerank 的分数计算和比较是依赖于模型 Embedding 本身的。所以这些知名的做检索的组织或公司，发表论文的方向主要是集中在 Embedding——如何让模型更好地表示和理解语义。至于 Rerank 基本都是在 Embedding 研究的基础上做 Cross-Encoder 或者 Bi-Encoder 即可。

所以下面我们会看到更多 Embedding 相关的模型。

BGE Rerank
#

BGE 相关的 Embedding 和 Rerank 模型目前搜到两篇论文，具体的模型可以直接去 Huggingface BAAI，源码地址如下：

FlagOpen/FlagEmbedding

Retrieval and Retrieval-augmented LLMs

Python

8881

646

下面简要概述一下他们的两篇论文：

[2023.12] Making Large Language Models A Better Foundation For Dense Retrieval

BGE Embedding 第一个版本主要是基于 LLaMA-2-7B (BASE) 模型做预训练和微调

这篇论文主要聚焦于仅解码器模型的痛点——模型在输出 Embedding 时侧重于当前本地和近未来的语义，忽视了句子整体上下文的语义。

提出了 LLaRA 架构：在预训练过程中让模型生成两个 Embedding。一个是 EBAR (Embedding-Based Auto-Encoding) 即原始句子的 Embedding 用于针对下一个标记的预测，另一个是 EBAR (Embedding-Based Auto-Regression)，即下一个句子的 Embedding 用于对句子级特征的预测，以此来增强模型对整体上下文的理解。

[2024.02] BGE M3-Embedding: Multi-Lingual, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation

BGE-M3 对比版本二做了全面的更新：将模型切换到了经过预训练的 XLM-RoBERTa，其次是支持了多语言、多种检索方式等等。最后还有一点是模型提出了自我知识蒸馏的训练框架，这点在这里就不展开说明了。

让我们回到 Rerank 本身，在做好了 Embedding 相关的设置后， Rerank 就比较简单了。

BGE-Rerank 是也是在 XLM-RoBERTa 的基础上进行微调来实现的，具体的方法使用了交叉编码器 Cross-Encoder，具体可以看源码：FlagOpen/FlagEmbedding - Reranker。

Jina Rerank
#

Jina 的 Embedding 和 Rerank 模型主要是基于 Bert 以及 Bert 模型相关变种的基础上进行预训练和微调。模型实例可以看 Huggingface jinaai，模型的源码貌似没有公开，这里我们主要看看相关论文：

[2023.07] Jina Embeddings: A Novel Set of High-Performance Sentence Embedding Models：
Jina 的第一版 Embedding 是基于 T5 模型 (*Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer.[2019.10]*) 进行预训练的，其次 Jina 团队设计了一个新的数据集专门用于训练 Embedding 模型。
[2023.10] Jina Embeddings 2: 8192-Token General-Purpose Text Embeddings for Long Documents：
Jina 推出的第二版 Embedding 为了实现长文本 8192 tokens，在 Bert 模型的基础上进行预训练，采用了线性偏差 Attention (ALiBi) 以增加模型输入的上下文信息；其次使用文本数据微调模型的 Embedding 能力；最后再次基础上添加反例微调，以此进一步提高模型能力。
[2024.09] jina-embeddings-v3: Multilingual Embeddings With Task LoRA：
Jina 推出的第三版本 Embedding 是一个比较新的模型。它是基于 XLM-RoBERTa 模型 (这里是借鉴了 BGE-M3) 进行训练和微调的，并引入了 FlashAttention 2 来提高计算效率。
比较有创新点的是，Jina 团队引入了五个用于特定任务的 LoRA 适配器，这五个适配器是独立训练的，具体有：① retrieval.passage：用于在查询文档检索任务中嵌入文档，主要是在非对称检索中的段落嵌入 ② retrieval.query：在查询文档检索任务中嵌入查询，主要是在非对称检索中的查询嵌入 ③ separation：对文档进行做聚类分析，用于聚类和重新排序的任务 ④ classification：文本分类，用于分类任务 ⑤ text-matching：语义文本相似度匹配，用于设计语义相似性的任务如对称检索。

回到 Rerank，目前 Jina 还没有发布基于 jina-embedding-v3 的 Rerank 模型 (估计快了)。

单就 jina-Reranker-v2 来讲，和 BGE-M3 一样采用了交叉编码器 Cross-Encoder。

GTE Rerank
#

GTE 相关的 Embedding 和 Rerank 是阿里巴巴团队发布的，阿里本身有大模型 Qwen 所以检索相关的嵌入和重排序基本也是基于 Qwen 来做训练和微调。具体的模型地址同样的可以看 Huggingface Alibaba-NL。

相关论文的话，主要是这篇：[2023.08] Towards General Text Embeddings with Multi-stage Contrastive Learning， ~~这里懒得写了~~ 。

为什么需要 Rerank
#

前面我们说到了目前很多 Rerank 是基于 Embedding 来实现的，Embedding 本身又是基于参数量相对小的大模型进行训练和微调的。

前面也介绍了很多 Embedding 模型一直在致力于如何更好地嵌入上下文语义，同时支持多语言、长文本等等等等。那么只用 Embedding 模型 + 向量检索算法也可以做到很好的效果，为什么还需要 Rerank？

从原理上来讲，个人认为区别主要在两点：

在 Embedding 向量嵌入层面，Rerank 模型目前大都使用的是 Cross-Encoder，即会把 query 和 document 拼接起来然后做 Embedding；而采用向量检索，是将 query 和 document 分别去做向量化，更类似上面所讲到的 Bi-Encoder
在分数设定方面，采用 Cross-Encoder 的 Rerank 模型会使用对数概率的方式，当然还有一些微调、适配器等等可能会影响最终的分数；而向量检索就是纯算法，有多种度量算法常用的主要还是余弦距离

从实际数据的角度说，这里引用 Cohere 的数据，总体来说 Rerank 还是会有一些提升的