BGE论文解读：如何炼成中文世界的通用Embedding Model

BGE论文解读：如何炼成中文世界的通用Embedding Model

在AI领域，Foundation Model是两个重要代表：Large Language Model和Embedding Model。前者通过text -> text的转换实现文本生成，后者则通过text -> embedding的转换将文本映射到向量空间，支持检索、聚类、分类等应用场景。本文将深入解读智源研究院的BGE论文，探讨如何构建中文世界的通用Embedding模型。

引言

Foundation Model有两个代表，一个是Large Language Model，另一个是Embedding Model。
前者聚焦文本空间，其形式化功能为text -> text；后者聚焦向量空间，其功能为text -> embedding。转为向量能做些什么呢？比较常见的使用场景包括
retrieval
（如检索知识库、检索Tool）、
clustering
（聚类）、
classification
（分类，通常需再接一层分类layer）等。
embedding model古已有之，但大多只擅长做某一特定任务；而要成为foundation model，必须满足通用这一关键特性，OpenAI是这方面的先行者[1]。
在中文世界，智源研究院的
BGE
是比较有名的开源embedding model，本篇博文即对BGE的论文 -- C-Pack: Packaged Resources To Advance General Chinese Embedding[2]进行解读。

方法论

BGE的目标是做中文世界的通用embedding模型。
通用，意味着用一个模型，支持所有的embedding使用场景，包括但不限于：retrieval、re-rank、clustering、classification、pair-classification等任务。
BGE从两个方面来达成这个目标：

• 数据方面，兼顾
  scale
  、
  diversity
  、
  quality
  这三个维度，这是通用embedding模型能训练出来的前提；
• 训练策略方面，论文使用3阶段训练策略，从
  pre-training
  到
  general-purpose fine-tuning
  再到
  task-specific fine-tuning
  ；前两个阶段是保证通用性的基石，最后一个阶段则在保持通用的基础上，进一步精进下游任务的效果。

数据方面

在训练数据方面，论文构建了大量的text pair数据，论文称之
C-MTP(Chinese Massive Text Pairs)
，数据量总计100M，涵盖多种任务，来自Wudao[3]等开源数据集，结合一些filter策略，同时达到scale、diversity、quality三个目标。
具体而言，
C-MTP
分成unlabeled和labeled两部分。
2. unlabeled数据，源于open web content和public Chinese dataset。前者包括
Wudao Corpora
、知乎、百科等数据，使用(title, passage)作为text pair；后者包括CSL、CMRC等公开数据集，这些数据集中pair结构天然存在，因此直接使用；同时，为了保证数据quality，使用Text2Vec-Chinese（预训练的中文embedding模型）[4]，过滤掉得分低于0.43的pair数据。最终数据量共计100M；
4. labeled数据，直接来自于下游任务的标注数据集，包括DuReader、mMARCO、NLI-Zh等，涵盖retrieval、ranking、similarity comparison等任务，数据量共计838K。
在测试数据方面，论文构建了中文世界的benchmark，称之为
C-MTEB(Chinese Massive Text Embedding Benchmark)
。

训练方面

简要来说，
pre-training
阶段在
Wudao Corpora
上进行，此阶段未在任何pair数据上训练，其目标是训练出更适合embedding任务的pre-trained model；
general-purpose fine-tuning
阶段在
C-MTP(unlabeled)
上进行，该阶段在100M的text pairs上训练，可以视作一种大规模的弱监督学习过程，可初步学习出通用embedding model；
最后的
task-specific fine-tuning
阶段，在
C-MTP(labeled)
上进行，通过在少而精的下游任务labeled data上微调，在保证通用性的同时，强化模型在具体任务上的表现。

训练细节

pre-training阶段

前文提到，这一阶段是为了学习出更适合embedding的pre-trained model。
论文采取了RetroMAE[5]的训练策略。其目标函数如下：
简单来说，就是先对text X进行随机Mask，然后进行encoding，再额外训练一个light-weight decoder（如单层transformer）进行重构。通过这一过程，强迫encoder学习到良好的embedding。
一个很自然的疑问是，这种方法比Bert要好吗？
作者对此进行了实验。其中
BGE-i w.o. pre-train
模型直接使用了Chinese-RoBERTa[6]，与本文的
BGE-i
模型对比，可以发现整体表现其实差不多，只是在retrieval任务上有较明显的提升。
retroMAE的提出也正是为了增强retrieval任务的表现。

general-purpose fine-tuning阶段

这一阶段的核心技术是对比学习，重点是：
2. 采用in-batch negative sample方法；
4. 使用大batch_size（论文使用的size为19200）。
这一阶段主打一个简单粗暴 --只要batch够大，在batch内就足以找到hard negative sample。
论文实验表明，
BGE-i
（仅经过general-purpose fine-tuning）的整体表现就已经可以超过
OpenAI-Ada-002
和
M3E（large）
了，细看一番，其中优势最大的是retrieval任务，这与其训练数据（
C-MTP-unlabel
）中包含大量web content中的（title，passage）有很大的关系。
着重优化retrieval任务的表现，应该是BGE团队的刻意为之；毕竟retrieval是embedding model最常见的使用场景。

task-specific fine-tuning阶段

这一阶段的难点在于：在任务间存在差异的情况下，如何更好地multi-task learning。
论文采取了两个关键技术：
2. instruction-based fine-tuning[7]。核心思路是将衡量
sim（x1，x2）
，转化为衡量
sim（instruction+x1，instruction+x2）
，这个instruction就是一段text prompt，用以说明domain、task等内容。例如在retrieval任务中，query侧加入的instruction为
；
4.