一文通透Text Embedding模型：从text2vec、openai-ada-002到m3e、bge

前言

本文已经是今年的第31篇大模型相关的技术文章了，如果说

半年之前写博客，更多是出于个人兴趣 + 读者需要
那自我司于23年Q3组建LLM项目团队之后，写博客就成了：个人兴趣 + 读者需要 + 项目需要
如此兼备三者，实在是写博客之幸运矣

我和我司更非常高兴通过博客、课程、内训、项目，与大家共同探讨如何把先进的大模型技术更好、更快的落地到各个行业的业务场景中，赋能千千万万公司的实际业务

而本文一开始是属于：因我司第三项目组「知识库问答项目」而起的此文《知识库问答LangChain+LLM的二次开发：商用时的典型问题及其改进方案》中的1.2节(该1.2节初稿来自我司LLM项目团队第三项目组的bingo)，但为把Text Embedding模型阐述的更为精准、全面，特把那部分的内容抽取出来，不断完善成此文

最终尽可能相比网上已有的其他资料都更细致化

第一部分衡量文本向量表示效果的榜单：MTEB、C-MTEB

1.2 《MTEB: Massive Text Embedding Benchmark(海量文本嵌入基准)》

判断哪些文本嵌入模型效果较好，通常需要一个评估指标来进行比较，《MTEB: Massive Text Embedding Benchmark(海量文本嵌入基准)》就是一个海量文本嵌入模型的评估基准

论文地址：https://arxiv.org/abs/2210.07316
MTEB包含8个语义向量任务，涵盖58个数据集和112种语言。通过在MTEB上对33个模型进行基准测试，建立了迄今为止最全面的文本嵌入基准。我们发现没有特定的文本嵌入方法在所有任务中都占主导地位。这表明该领域尚未集中在一个通用的文本嵌入方法上，并将其扩展到足以在所有嵌入任务上提供最先进的结果
github地址：https://github.com/embeddings-benchmark/mteb#leaderboard

榜单地址：https://huggingface.co/spaces/mteb/leaderboard

1.2 中文海量文本embedding任务排行榜：C-MTEB

从Chinese Massive Text Embedding Benchmark中可以看到目前最新的针对中文海量文本embedding的各项任务的排行榜，针对不同的任务场景均有单独的排行榜。

任务榜单包括：

Retrieval
STS
PairClassification
Classification
Reranking
Clustering

其中，在本地知识库任务中，主要是根据问题query的embedding表示到向量数据库中检索相似的本地知识文本片段。因此，该场景主要是Retrieval检索任务。检索任务榜单如下：

目前检索任务榜单下效果最好的是bge系列的bge-large-zh模型，langchain-chatchat项目中默认的m3e-base也处于比较靠前的位置

第二部分 text-embedding-ada-002

2.1 模型简介

text-embedding-ada-002是OpenAI于2022年12月提供的一个embedding模型，但需要调用接口付费使用。其具有如下特点：

统一能力：OpenAI通过将五个独立的模型(文本相似性、文本搜索-查询、文本搜索-文档、代码搜索-文本和代码搜索-代码)合并为一个新的模型
在一系列不同的文本搜索、句子相似性和代码搜索基准中，这个单一的表述比以前的嵌入模型表现得更好
上下文：上下文长度为8192，使得它在处理长文档时更加方便
嵌入尺寸：只有1536个维度，是davinci-001嵌入尺寸的八分之一，使新的嵌入在处理矢量数据库时更具成本效益

2.2 模型使用

以下是OpenAI官方文档中给出的用于文本搜索的代码实例

from openai.embeddings_utils import get_embedding, cosine_similarity

def search_reviews(df, product_description, n=3, pprint=True):
  embedding = get_embedding(product_description, model='text-embedding-ada-002')
  df['similarities'] = df.ada_embedding.apply(lambda x: cosine_similarity(x, embedding))
  res = df.sort_values('similarities', ascending=False).head(n)
  return res

res = search_reviews(df, 'delicious beans', n=3)

第三部分 m3e模型

3.1 m3e模型简介

M3E(Moka Massive Mixed Embedding，分为m3e-small、m3e-base、m3e-large，m3e GitHub地址：GitHub - wangyingdong/m3e-base)，其

使用in-batch负采样的对比学习的方式在句对数据集进行训练，为了保证in-batch负采样的效果，使用A100来最大化batch-size，并在共计2200W+的句对数据集(包含中文百科，金融，医疗，法律，新闻，学术等多个领域)上训练了 1 epoch
使用了指令数据集，M3E 使用了300W+的指令微调数据集，这使得 M3E 对文本编码的时候可以遵从指令，这部分的工作主要被启发于 instructor-embedding
基础模型，M3E 使用 Roberta 系列模型进行训练，目前提供 small 和 base 两个版本
此文《知识库问答LangChain+LLM的二次开发：商用时的典型问题及其改进方案》中的langchain-chatchat便默认用的m3e-base

3.1.1 m3e与openai text-embedding-ada-002

以下是m3e的一些重大更新

2023.06.08，添加检索任务的评测结果，在 T2Ranking 1W 中文数据集上，m3e-base 在 ndcg@10 上达到了 0.8004，超过了 openai-ada-002 的 0.7786
见下图s2p ndcg@10那一列(其中s2p, 即 sentence to passage ，代表了异质文本之间的嵌入能力，适用任务：文本检索，GPT 记忆模块等)
2023.06.07，添加文本分类任务的评测结果，在 6 种文本分类数据集上，m3e-base 在 accuracy 上达到了 0.6157(至于m3e-large则是0.6231)，超过了 openai-ada-002 的 0.5956
见下图s2s ACC那一列(其中s2s, 即 sentence to sentence ，代表了同质文本之间的嵌入能力，适用任务：文本相似度，重复问题检测，文本分类等)

此外，m3e团队建议

如果使用场景主要是中文，少量英文的情况，建议使用 m3e 系列的模型
多语言使用场景，并且不介意数据隐私的话，作者团队建议使用 openai text-embedding-ada-002
代码检索场景，推荐使用 openai text-embedding-ada-002
文本检索场景，请使用具备文本检索能力的模型，只在 S2S 上训练的文本嵌入模型，没有办法完成文本检索任务

3.2 m3e模型微调

微调脚本：
m3e是使用uniem脚本进行微调

from datasets import load_dataset

from uniem.finetuner import FineTuner

dataset = load_dataset('shibing624/nli_zh', 'STS-B')
# 指定训练的模型为 m3e-small
finetuner = FineTuner.from_pretrained('moka-ai/m3e-small', dataset=dataset)
finetuner.run(epochs=3)

详细教程暂放在「大模型项目开发线上营」中，至于本文后续更新

第四部分 bge模型

4.1 bge模型的简介

2023年8月2日，北京智源人工智能研究院发布的中英文语义向量模型BGE(hf地址：https://huggingface.co/BAAI/bge-large-zh，GitHub地址：https://github.com/FlagOpen/FlagEmbedding/blob/master/README_zh.md)，以下是BGE的技术亮点

高效预训练和大规模文本微调；
在两个大规模语料集上采用了RetroMAE预训练算法，进一步增强了模型的语义表征能力；
通过负采样和难负样例挖掘，增强了语义向量的判别力；
借鉴Instruction Tuning的策略，增强了在多任务场景下的通用能力

4.1.1 RetroMAE的预训练步骤

目前主流的语言模型的预训练任务都是token级别的，比如MLM或者Seq2Seq，但是这种训练任务难以让模型获得一个高质量的基于句子级别的句向量，这限制了语言模型在检索任务上的潜力。针对这个弊端，目前有两者针对检索模型的预训练策略

第一种是self-contrastive learning，这种方式往往受限于数据增强的质量，并且需要采用非常庞大数量的的负样本
另一种基于anto-encoding，一种自重建方法，不受数据增强跟负样本采样策略的影响，基于这种方法的模型性能好坏有两个关键因素
其一是重建任务必须要对编码质量有足够的要求，其二是训练数据需要被充分利用到

基于此，研究人员提出了RetraoMAE(RetroMAE论文：https://arxiv.org/abs/2205.12035)，它包括两个模块，其一是一个类似于BERT的编码器，用于生成句向量，其二是一个一层transformer的解码器，用于重建句子，如下图所示

4.1.1.1 编码Encoding

所谓编码，即Mask(EN)掉一小部分token然后通过BERT编码得到句子嵌入sentence embedding $h_x$ ，具体步骤如下

给定一个句子输入 $X$ ：Norwegian forest cat is a breed of dom-estic cat originating in northern Europe
随机Mask(EN)掉其中一小部分token后得到 $X_{enc}$ ：[M] forest cat is a breed of [M] cat originating in [M] Europe
这里通常会采用一定的mask比例(15%～30%)，从而能保留句子原本大部分的信息
然后利用类似BERT的编码器 $\Phi^{e n c}(\cdot)$ 对其进行编码，得到对应的的句子嵌入 $\mathbf{h}_{\tilde{X}}$ 「一般将[CLS]位置最后一层的隐状态作为句子嵌入」，如下公式所示
$\mathbf{h}_{\tilde{X}} \leftarrow \Phi_{e n c}\left(\tilde{X}_{e n c}\right)$
We apply a BERT like encoder with 12 layers and768 hidden-dimensions, which helps to capture thein-depth semantics of the sentence. Following the common practice, we select the [CLS] token’s finalhidden state as the sentence embedding.

4.1.1.2 解码Decoding

所谓解码，即Mask(DE)很大一部分token然后结合句子嵌入sentence embedding $\mathbf{h}_{\tilde{X}}$ ，让解码器重构原始句子

具体而言，即是联合以下两个部分，好让解码器在该两部分的基础上重构原始句子

利用Mask(DE)后的文本输入 $\tilde{X}_{d e c}$ ：[MI [M] cat is MI [M} of dom-estic [M] [M] in northern [M]
(这里采取了比encoder部分更加激进的mask比例，比如50%～70%)
与上一节encoder生成的句子嵌入(sentence embedding)
论文中对这一步骤的英文阐述是：The masked input is joined with the sentence embedding, based on which the original sentence is reconstructed by the decoder.

有个细节是，这里的Mask(DE)输入带上位置嵌入了，即句子嵌入 $\mathbf{h}_{\tilde{X}}$ 和带有位置的掩码输入被组合成以下序列
$\mathbf{H}_{\tilde{X}_{d e c}} \leftarrow\left[\mathbf{h}_{\tilde{X}}, \mathbf{e}_{x_{1}}+\mathbf{p}_{1}, \ldots, \mathbf{e}_{x_{N}}+\mathbf{p}_{N}\right]$
其中， $e_{x_i}$ 表示 $x_i$ 的嵌入，在 $x_i$ 的基础上增加了一个额外的位置嵌入 $p_i$

接下来，通过优化以下目标，学习解码器 $\Phi_{d e c}$ 来重构原始句子 $X$
$\mathcal{L}_{d e c}=\sum_{x_{i} \in \text { masked }} \mathrm{CE}\left(x_{i} \mid \Phi_{\text {dec }}\left(\mathbf{H}_{\tilde{X}_{d e c}}\right)\right)$
其中， $CE$ 是交叉熵损失

由于在解码器部分采用了极其简单的网络结构跟非常激进的mask比例，从而使得解码任务变得极具挑战性，迫使encoder去生成高质量的句向量才能最终准确地完成原文本重建

4.1.1.3 增强解码Enhanced Decoding

前面提及的解码策略有一种缺陷，就是训练信号只来源于被mask掉的token，而且每个mask掉的token都是基于同一个上下文重建的。于是研究人员提出了一种新的解码方法：Enhanced Decoding，具体做法如下

a) 首先生成两个不同的输入流H1(query)跟H2(context)
其中 $\mathbf{h}_{\tilde{X}}$ 是句子嵌入， $e_{x_i}$ 是标记嵌入(在这个地方没有标记被掩码)， $p_i$ 是位置嵌入
相当于
$H_1$ 是sentence embedding + Position embedding
$H_2$ 为sentence embedding和token embedding + position embedding
b) 通过attention机制得到新的输出 $A$
这里的 $M$ 是一个mask矩阵，第 $i$ 个token所能看得到的其他token是通过抽样的方式决定的(当然要确保看不到自身token，而且都要看得见第一个token，也就是encoder所产出CLS句向量的信息)

其中一个和常规decoder不一样的地方是，H1作为Q，H2作为KV

H1中的每个token embedding去H2中查找比较重要的上下文：包括H2被采样到的 token，以及初始token embedding都能看到[这里的初始embedding就是sentence embedding]，至于对角线上的因代表的各自自身，故看不到

为方便大家更好、更快的理解，我再举个例子，比如：

P0上的 $x_0$ 能看见P1上的 $x_1$ (x0、x1)
P1上的 $x_1$ 能看见P0、P2上的 $x_0$ 、 $x_2$ (x0、x1、x2)
P2上的 $x_2$ 能看见P0、P1上的 $x_0$ 、 $x_1$ (x0、x1、x2)
P3上的 $x_3$ 能看见P0、P4上的 $x_0$ 、 $x_4$ (x0、_ 、 _、x3、x4)
P4上的 $x_4$ 能看见P0、P3上的 $x_0$ 、 $x_3$ (x0、_ 、 _、x3、x4)
之后，每个token xi基于对矩阵M的第i行可见的上下文进行重构(each token xi is recon-structed based on the context which are visible tothe i-th row of matrix M)，该矩阵即如下所示
$\mathbf{M}_{i j}=\left\{\begin{array}{l} 0, \quad x_{j} \in s\left(X_{\neq i}\right), \text { or } j_{\mid i \neq 0}=0 \\ -\infty, \text { otherwise. } \end{array}\right.$
主对角线位置填充为-∞(因为其代表自身，故不可见)，可见上下文的位置填充为0，代表可见

c)最终利用attention后的输出A跟H1一起过FNN(即resnet)去重建原文本，这里重建的目标不仅仅是被mask掉的token，而是全部token
$\mathcal{L}_{d e c}=\sum_{x_{i} \in X} \mathrm{CE}\left(x_{i} \mid \mathbf{A}, \mathbf{H}_{1}\right)$

最终RetroMAE的损失由两部分相加得到，其一是encoder部分的MLM损失，其二是deocder部分自重建的交叉熵损失

最后，再总结一下RetroMAE 预训练步骤

(A)编码阶段：将输入进行一定比例的mask操作，并编码为句子嵌入(绿色矩形)
(A) Encoding: the input is moderately masked and encoded as the sentence embedding (the green rectangle)
(B)解码阶段：对输入使用很高比例的mask操作，并与句子嵌入连接以恢复被mask的部分(阴影符号)
(B) Decoding: the input is aggressively masked, and joined with the sentence embedding to reconstruct the masked tokens (the shadowed tokens).
(C)增强编码阶段：基于每行的句子嵌入和可见上下文来重建所有输入符号；主对角线位置填充为-∞(灰色，因为其代表自身，故不可见)，可见上下文的位置填充为0(蓝色)
(C) Enhanced encoding: all input tokens are reconstructed based on the sentence embedding and the visible context in each row (defined in Eq. 7); the main diagonal positions are filled with −∞ (grey), and positions for the visible context are filled with 0 (blue).

4.2 bge模型的微调

微调脚本：https://github.com/FlagOpen/FlagEmbedding/tree/master/examples/finetune

数据格式

{"query": str, "pos": List[str], "neg":List[str]}

难负样本挖掘
难负样本是一种广泛使用的提高句子嵌入质量的方法。可以按照以下方法挖掘难负样本

python -m FlagEmbedding.baai_general_embedding.finetune.hn_mine \
--model_name_or_path BAAI/bge-base-en-v1.5 \
--input_file toy_finetune_data.jsonl \
--output_file toy_finetune_data_minedHN.jsonl \
--range_for_sampling 2-200 \
--use_gpu_for_searching

训练

python -m FlagEmbedding.baai_general_embedding.finetune.hn_mine \
--model_name_or_path BAAI/bge-base-en-v1.5 \
--input_file toy_finetune_data.jsonl \
--output_file toy_finetune_data_minedHN.jsonl \
--range_for_sampling 2-200 \
--use_gpu_for_searching