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

前言

本文一开始是属于此文《知识库问答LangChain+LLM的二次开发：商用时的典型问题及其改进方案》的1.2节，但为把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

// 待更

第三部分 m3e模型

3.1 m3e模型简介

M3E(Moka Massive Mixed Embedding，m3e-base模型下载地址：https://huggingface.co/moka-ai/m3e-base，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，超过了 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模型的简介

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的解码器，用于重建句子，如下图所示

编码Encoding
给定一个句子输入X：Norwegian forest cat is a breed of dom-estic cat originating in northern Europe
$\rightarrow$ 随机Mask(EN)掉其中一小部分token后得到X_enc：[M] forest cat is a breed of [M] cat originating in [M] Europe
这里通常会采用一定的mask比例(15%～30%)，从而能保留句子原本大部分的信息
$\rightarrow$ 然后利用类似BERT的编码器对其进行编码，得到对应的的句子嵌入 $h_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.
解码Decoding
利用Mask(DE)后的文本输入(但这里采取了比encoder部分更加激进的mask比例，比如50%～70%)：[MI [M] cat is MI [M} of dom-estic [M] [M] in northern [M]
与encoder生成的句子嵌入(sentence embedding)连接，解码器重构原始句子
The masked input is joined with the sentence embedding, based on which the original sentence is reconstructed by the decoder.

且句子嵌入和掩码输入被组合成以下序列
$\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)$

由于在解码器部分采用了极其简单的网络结构跟非常激进的mask比例，从而使得解码任务变得极具挑战性，迫使encoder去生成高质量的句向量才能最终准确地完成原文本重建
增强解码Enhanced Decoding
前面提及的解码策略有一种缺陷，就是训练信号只来源于被mask掉的token，而且每个mask掉的token都是基于同一个上下文重建的。于是研究人员提出了一种新的解码方法，Enhanced Decoding，具体做法如下
a) 首先生成两个不同的输入流H1(query)跟H2(context)
其中 $\mathbf{h}_{\tilde{X}}$ 是句子嵌入， $e_{x_i}$ 是标记嵌入(在这个地方没有标记被掩码)， $p_i$ 是位置嵌入

b) 通过attention机制得到新的输出 $A$ ，这里的 $M$ 是一个mask矩阵，第 $i$ 个token所能看得到的其他token是通过抽样的方式决定的(当然要确保看不到自身token，而且都要看得见第一个token，也就是encoder所产出CLS句向量的信息)
c)最终利用A跟H1去重建原文本，这里重建的目标不仅仅是被mask掉的token，而是全部token
最终RetroMAE的损失由两部分相加得到，其一是encoder部分的MLM损失，其二是deocder部分自重建的交叉熵损失

总结一下RetroMAE 预训练步骤，如下图所示

(A)编码阶段：将输入进行一定比例的掩码操作，并编码为句子嵌入(绿色矩形)
(A) Encoding: the input is moderately masked and encoded as the sentence embedding (the green rectangle)
(B)解码阶段：对输入使用很高比例的掩码操作，并与句子嵌入连接以恢复掩码符号(阴影符号)
(B) Decoding: the input is aggressively masked, and joined with the sentence embedding to reconstruct the masked tokens (the shadowed tokens).
(C)增强编码阶段：基于每行的句子嵌入和可见上下文(在Eq. 7中定义)来重建所有输入符号；主对角线位置填充为-∞(灰色)，可见上下文的位置填充为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.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