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

引言

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两部分。

1. 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；
2. 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阶段

这一阶段的核心技术是对比学习，重点是：

1. 采用in-batch negative sample方法；
2. 使用大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。

论文采取了两个关键技术：

1. instruction-based fine-tuning[7]。核心思路是将衡量 sim（x1，x2），转化为衡量 sim（instruction+x1，instruction+x2），这个instruction就是一段text prompt，用以说明domain、task等内容。例如在retrieval任务中，query侧加入的instruction为 为这个句子生成表示以用于检索相关文章：；
2. hard negative sampling 。在训练过程中，采取ANN-style sampling strategy[8]，从该任务的corpus中全局性 地采样出一个embedding表征最接近的hard negative sample 。

论文实验表明， BGE-f（经过task-specific fine-tuning之后的最终模型），在retrieval、STS、pair-classification、re-rank上 明显好于 BGE-i模型，其他任务也 几乎没有效果损失 。

这充分说明了论文所采取的fine-tuning技术的有效性 。

可以认为，这也是BGE团队做出的fine-tuning示范。

BGE的效果如何

BGE所采取的模型结构类似于BERT，使用其最后一层中 CLS token的hidden state作为embedding。

BGE模型有3个不同大小的版本，其中small版参数量为 24M 、base版参数量为 102M 、large版参数量为 326M 。

论文在 C-MTEB上与众多embedding模型进行了对比，结果如下。

可以看到， BGE-base和BGE-large基本是一骑绝尘 ，几乎在每一个任务上的效果都明显更好， 即使是BGE-small，也几乎能达到SOTA的效果 。

总结

本文对BGE的论文进行了简要解读，指出其在数据侧、训练侧所采取的方法论，同时对其三阶段训练过程的技术细节进行了介绍。

一个很明显的感受： BGE是一个产品 ，它做到的是集各家之所长（数据、训练技术等），产品的诞生便是它的创新之处。从论文实验来看，在中文的各项任务上，就算是 small-size BGE，也可以战胜 OpenAI-Ada-002，而且BGE是开源的 -- 这意味着每个人都可以对它再次进行fine-tuning，以进一步提升效果。这的确是中文开发者的福音。

关注笔者

欢迎关注笔者，了解有趣、有用的NLP知识：）

参考资料

[1] Text and Code Embeddings by Contrastive Pre-Training: https://cdn.openai.com/papers/Text\_and\_Code\_Embeddings\_by\_Contrastive\_Pre\_Training.pdf

[2] C-Pack: Packaged Resources To Advance General Chinese Embedding: https://arxiv.org/abs/2309.07597

[3] Wudao Corpus: https://data.baai.ac.cn/details/WuDaoCorporaText

[4] Text2Vec-Chinese: https://huggingface.co/GanymedeNil

[5] RetroMAE: https://aclanthology.org/2022.emnlp-main.35/

[6] Chinese-RoBERTa: https://huggingface.co/hfl/chinese-roberta-wwm-ext-large

[7] One Embedder, Any Task: Instruction-Finetuned Text Embeddings: https://aclanthology.org/2023.findings-acl.71/

[8] Approximate Nearest Neighbor Negative Contrastive Learning for Dense Text Retrieval: https://openreview.net/forum?id=zeFrfgyZln