文本向量模型BGE与BGE-M3

BGE模型

BGE模型对应的技术报告为《C-Pack: Packaged Resources To Advance General Chinese Embedding》

训练数据

为了训练BGE向量模型，构建了C-MTP数据集，它包括了用来训练文本向量的文本对数据（问答对、两个同义句子、相同主题的两个文档等），为了保证文章向量的泛化性，数据集需要同时满足大规模和多样性，C-MTP包括两个部分：未标注数据C-MTP (unlabeled)和标注数据C-MTP (labeled)。

• C-MTP (unlabeled)：数据主要来源是开放网络语料库如Wudao corpus，对于每一篇文章抽取(title, passage)来构成一个文本对；也按同样的处理方式从知乎、百科等抽取了文本对。除了网络开放语料之外，还从以下公开中文数据集中抽取了文本对：
  • CSL (scientific literature)
  • Amazon-Review-Zh (reviews)
  • Wiki Atomic Edits (paraphrases)
  • CMRC (machine reading comprehension)
  • XLSUM-Zh (summarization)

​ 直接抽取的文本对可能质量不高，使用Text2VecChinese给文本对打分，只保留分数在0.43上的文本对，数据集最终文本对的大小为100M

• C-MTP (labeled)：标注文本对来自以下标注数据集，一共838465条：
  • T2-Ranking
  • DuReader
  • mMARCO
  • CMedQA-v2
  • multi-cpr
  • NLI-Zh3
  • cmnli
  • ocnli

训练流程

BGE模型基于BERT-like架构，使用最后一层的特殊token [CLS]作为文本向量，有三个不同规模的模型组成：

• large (326M参数)
• base (102M参数)
• small (24M 参数)

BGE的训练分为三个部分：1)预训练，2)通用微调（用C-MTP (unlabeled)进行对比学习），3)任务相关微调（用C-MTP (labeled)进行多任务微调学习）。

1. 预训练，使用Wudao语料来训练模型，使用RetroMAE中提出的自动编码方式(MAE-stype)来进行预训练。

2. 通用微调，将预训练好的模型在C-MTP (unlabeled)数据集上通过对比学习微调。

• $min.{\sum }_{(p,q)}-\log \frac{e^{<e_p,e_q>}/\tau }{e^{<e_p,e_q>/\tau }+{\sum }_{Q^{\prime }}{e}^{<e_p,e_{q^{\prime }}>/\tau }}$，式中p和q是文本对，$q^{\prime }\in Q^{\prime }$是负样本，$\tau$是温度。
• 负样本仅使用in-batch negative samples，batch size 高达19200。
• 训练时使用gradient checkpointing 和 cross-device embedding sharing的组合策略来使得batch size可以很大

3. 任务相关微调，将前一步训练好的向量模型在C-MTP (labeled)数据集上进一步微调。为了协调不同任务，采用如下两个策略：

• 基于指令的微调，对于每一个文本对(p,q)，将一个任务相关的指令$I_t$添加到query q上：$q^{\prime }\leftarrow q+I_t$，指令是一个用来描述任务的prompt，比如"search relevant passages for the query".
• 对于每一个文本对(p,q)，挖掘了一个hard negative sample $q^{\prime }$，挖掘方法是基于ANN挖掘的。

消融实验表明训练的各个阶段对于模型效果提升都有帮助（w.o.Instruct是在第三阶段训练时不使用指令微调，BGE-i是指经过第二阶段通用微调的模型，BGE-i w.o. pre-train是不使用RetroMAE预训练的模型，BGE-f经过全部训练流程的模型）

也做了实验来说明大batch-size对于模型性能的影响，在大多数任务上大batch size对于embedding模型是有益处的，尤其是检索任务效果提升特别明显。

BGE-M3模型

BGE-M3对应的技术报告为《BGE M3-Embedding: Multi-Lingual, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation》

BGE-M3的M3是Multi-Linguality,Multi-Functionality, and Multi-Granularity的简称，也就是模型支持多语言、多种能力（稠密检索、稀疏检索、多向量检索），多粒度（最高支持8192长度的输入文本）。

训练数据

BGE-M3的数据构造分为三个部分：从未标注语料得到的无监督数据、从标注语料得到的微调数据、合成微调数据。

• 无监督数据共1.2 billion文本对，包括了194种语言和2655跨语言对。主要从不同的多语言语料库提取titlebody, title-abstract, instruction-output等，数据来源有：

  • 多语言语料库：Wikipedia、S2ORC、xP3、mC4、CC-News
  • MTP（BGE模型的数据集）
  • 翻译相关数据集：NLLB、CCMatrix

• 从标注语料得到的微调数据，数据来源有：

  • 英文包括8个数据集：HotpotQA，TriviaQA，NQ，MS MARCO，COLIEE，PubMedQA， SQuAD， NLI。
  • 中文包括7个数据集：DuReader，mMARCO-ZH， $T^2$-Ranking， LawGPT，CMedQAv2，NLI-zh，LeCaRDv2。
  • 其他语言：MIRACL, Mr.TyDi

• 合成数据(MultiLongDoc)，从Wiki和MC4数据集中采样了长文章，并从其中随机选择段落，由GPT-5基于这些段落生成问题。生成的问题和段落组成的文本对最终作为微调数据的一部分。

  • GPT 3.5使用的prompt为“You are a curious AI assistant, please generate one specific and valuable question based on the following text. The generated question should revolve around the core content of this text, and avoid using pronouns (e.g., “this”). Note that you should generate only one question, without including additional content:”

BGE-M3混合检索实现

BGE-M3的Multi-Functionality是指实现混合检索，具体如下：

• 稠密检索(Dense retrieval)，输入query q经过文本编码器的隐状态层${\mathbf{H}}_q$得到，取"[CLS]"的归一化隐状态作为向量：$e_q=norm({\mathbf{H}}_q[0])$，类似地passage p的向量：$e_p=norm({\mathbf{H}}_p[0])$。 query和passage的相关分数由两个向量$e_q$和$e_p$的内积表示：$s_{dense}\leftarrow ⟨e_p,e_q⟩$。
• 词性检索(Lexical Retrieval)，输入query的每一元素term t（即 token），term权重为$w_{q_t}\leftarrow Relu({\mathbf{W}}_{lex}^T{\mathbf{H}}_q[i])$，式中的${\mathbf{W}}_{lex}\in R^{d\times 1}$是映射隐状态到浮点数的矩阵（如果一个term t在query中出现多次，只保留其最大值）。按照同样的方式计算passage的term权重。query和passage的相关分数由两者共现项（记作$q\cap p$）计算得到：$s_{lex}\leftarrow {\sum }_{t\in q\cap p}(w_{q_t}\ast w_{p_t})$。
• 多向量检索(Multi-Vector Retrieval)，多向量检索使用query和passage的整个输出向量：$E_q=norm({\mathbf{W}}_{\mathbf{mul}}^{\mathbf{T}}{\mathbf{H}}_q)$， $E_p=norm({\mathbf{W}}_{\mathbf{mul}}^{\mathbf{T}}{\mathbf{H}}_p)$，式中的${\mathbf{W}}_{mul}\in R^{d\times d}$是可学习映射矩阵。query和passage的相关分数为：$s_{mul}\leftarrow \frac{1}{N}{\sum }_{i=1}^{N}ma{x}_{j=1}^M{E}_q[i]\cdot E_p^T[j]$， N和M分别是query和passage的长度。

因为向量模型的多功能能力，所以可以实现混合检索过程，首先候选结果由每种方法单独得到（多向量检索成本很高，可考虑省略），然后最后检索结果根据各相关分数之和来排序：$s_{rank}\leftarrow s_{dense}+s_{lex}+s_{mul}$。

Self-Knowledge Distillation

向量模型训练是为了将正样本从负样本中区分开，损失函数为如下式所表示的InfoNCE损失函数，式中$p^{\ast }$和$P^{\prime }$是query q对应的正样本及负样本。$s(\cdot )$ 是 { s d e n s e ( ⋅ ) , s l e x ( ⋅ ) , s m u l ( ⋅ ) } \{s_{dense}(\cdot) , s_{lex}(\cdot) , s_{mul}(\cdot) \} {sdense​(⋅),slex​(⋅),smul​(⋅)}中的任意一个。
L = − log ⁡ e x p ( s ( q , p ∗ ) / τ ) ∑ p ∈ { p ∗ , P ′ } e x p ( s ( q , p ) / τ ) \mathcal{L} = -\log \frac{exp(s(q, p^*)/\tau)}{\sum_{p \in \{p^*, P^{\prime}\}} exp(s(q,p)/\tau) } L=−log∑p∈{p∗,P′}​exp(s(q,p)/τ)exp(s(q,p∗)/τ)​
BGE-M3中的三种不同检索方法的目标函数可能会彼此冲突，解决办法是使用Self-Knowledge Distillation。

1. 三种检索方法融合的最简单方法是直接求和：$s_{inter}\leftarrow s_{dense}+s_{lex}+s_{mul}$。

2. 对于每一种检索方法可以使用融合分数$s_{inter}$作为teacher，因此其损失函数可修改成：${\mathcal{L}}_{\ast }^{\prime }\leftarrow -p(s_{inter})\ast \log p(s_{\ast })$。式中的$p(\cdot )$是softmax激活函数，$s_{\ast }$是 { s d e n s e ( ⋅ ) , s l e x ( ⋅ ) , s m u l ( ⋅ ) } \{s_{dense}(\cdot) , s_{lex}(\cdot) , s_{mul}(\cdot) \} {sdense​(⋅),slex​(⋅),smul​(⋅)}中的任意一个。

3. 将前一步的损失函数进行融合与归一化：${\mathcal{L}}^{\prime }\leftarrow ({\mathcal{L}}_{dense}^{\prime }+{\mathcal{L}}_{lex}^{\prime }+{\mathcal{L}}_{mul}^{\prime })/3$

4. Self-Knowledge Distillation最终的损失函数为$\mathcal{L}$和${\mathcal{L}}^{\prime }$的线性组合：${\mathcal{L}}_{final}\leftarrow \mathcal{L}+{\mathcal{L}}^{\prime }$

BGE-M3训练过程

像BGE一样，BGE-M3也是多阶段训练：

1. 经过RetroMAE 方法继续预训练的XLM-RoBERTa作为文本编码器。扩充了XLM-RoBERTa的最大位置向量到8192。数据来源为PIle、Wudao、mC4数据集，从中采样了184M数据，覆盖了105种语言。

2. 使用无监督数据用对比学习方式预训练，这一阶段只有稠密检索向量被训练。

3. 应用Self-Knowledge Distillation技术对向量模型继续微调，标注语料和合成数据都在这个阶段使用，并使用ANN方法挖掘了难负例。

为了保证训练过程中的大batch size 和长本文序列训练，优化了训练过程：

• 根据长度对文本数据进行分组，并从每组中采样数据，确保一个batch内文本长度相对相似，从而减少填充。同时，数据程序会首先从组内采样足够的数据，然后分配给各机器（random seed保持固定），保证不同机器的计算开销尽可能相近。
• 为了减少文本建模时的显存消耗，将一批数据分成多个小批。对于每个小批，我们利用模型编码文本，收集输出的向量同时丢弃所有前向传播中的中间状态，最后汇总向量计算损失。这个方法可以显著增加batch size。
• 不同GPU的embedding向量被广播，允许每台机器在分布式环境里获取到所有向量，可以扩展in-batch 负样本的规模

考虑到用户可能没有足够的资源或者数据来训练长文本模型，提出了一个MCLS策略使得无需训练就能增强模型的长文本能力。这个策略在推理时应用，使用多个CLS token来联合捕获长文本的语义。该策略为每个固定数量的令牌插入一个CLS token（论文使用的是256个token），每个CLS token可以从相邻的令牌获取语义信息。最后通过对所有CLS token的最后隐藏状态求平均值来获得最终的文本向量。

参考资料

1. Xiao, Shitao, Zheng Liu, Peitian Zhang, and Niklas Muennighof. 2023. “C-Pack: Packaged Resources To Advance General Chinese Embedding,” September.

2. Chen, Jianlv, Shitao Xiao, Peitian Zhang, Kun Luo, Defu Lian, and Zheng Liu. n.d. “BGE M3-Embedding: Multi-Lingual, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation.”

3. FlagEmbedding github