项目地址:https://huggingface.co/BAAI/bge-m3
首先说明下为什么需要text embedding?
Text Embedding 是一种将文本数据映射到高维向量空间的技术,这些向量能够捕捉文本的语义信息。通过将文本嵌入到向量空间中,我们可以利用向量之间的距离或相似性来衡量文本之间的语义相关性。
在信息检索领域,Text Embedding 的重要性尤为突出。传统的基于关键词匹配的检索方法往往无法准确理解用户的真实意图,也无法处理语义相似但用词不同的查询。而通过 Text Embedding,我们可以将查询和文档都嵌入到同一个语义空间中,从而能够更准确地检索出与用户查询语义相关的文档,而不仅仅是包含相同关键词的文档。
此外,随着多语言内容的不断增加,跨语言检索的需求也日益增长。Text Embedding 能够将不同语言的文本映射到一个统一的语义空间中,使得跨语言检索成为可能。
智源发布的BGE-M3在多语言、多功能和多粒度方面展现出前所未有的通用性,能够支持超过100种语言的语义检索,并且可以处理从句子到长文档的各种输入。下面来详细看下:
1、方法介绍
M3-Embedding 实现了三方面的多功能性,支持多种语言,并能够处理不同粒度的输入数据。此外,它还统一了文本嵌入的常见检索功能。形式上,给定任意语言 x 中的查询 q,它能够从语料库中检索出语言 y 中的文档
。
数据整理
BGE-M3-Embedding 需要一个大规模且多样化的多语言数据集。为了构建这样的数据集,从三个来源进行了全面的数据收集:
- 无监督数据:来源包括 Wikipedia、S2ORC、xP3、mC4、CC-News 和 MTP 等多语言语料库。从这些语料库中提取丰富的语义结构,例如标题-正文、标题-摘要、指令-输出等。为了学习跨语言语义匹配的统一嵌入空间,引入了平行句子数据,来自 NLLB 和 CCMatrix 两个翻译数据集。总共整理出1.2 亿个文本对,涵盖 194 种语言和 2655 种跨语言对应关系。
- 微调数据:
英语:整合了 8 个数据集,包括 HotpotQA、TriviaQA、NQ、MS MARCO、COLIEE、PubMedQA、SQuAD 和 SimCSE 的 NLI 数据。
中文:整合了 7 个数据集,包括 DuReader、mMARCO-ZH、、LawGPT、CMedQAv2、
和 LeCaRDv2。
其他语言:利用 Mr. TyDi 和 MIRACL 的训练数据。 - 合成数据:为缓解长文档检索任务的短缺,生成了额外的多语言微调数据(称为 MultiLongDoc)。从 Wikipedia、Wudao 和 mC4 数据集中抽取长篇文章,并从中随机选择段落。使用 GPT-3.5 根据这些段落生成问题,生成的问题和抽取的文章构成新的文本对。
混合检索
BGE-M3-Embedding 统一密集检索、词汇(稀疏)检索和多向量检索的具体实现如下:
- 密集检索:输入查询q通过文本编码器转换为隐藏状态
,使用特殊标记 "[CLS]" 的归一化隐藏状态表示查询:
。同样,段落p的嵌入为
。查询和段落之间的相关性得分通过内积计算:
。
- 词汇检索:输出嵌入用于估计每个词项的重要性,以促进词汇检索。对于查询中的每个词项t,词项权重计算为
,其中
是将隐藏状态映射到浮点数的矩阵。如果词项t在查询中出现多次,只保留其最大权重。查询和段落之间的相关性得分通过查询和段落中共存词项的联合重要性计算:
。
- 多向量检索:作为密集检索的扩展,多向量方法利用整个输出嵌入来表示查询和段落:
,其中
是可学习的投影矩阵。使用延迟交互计算细粒度的相关性得分:
,其中N和M分别是查询和段落的长度。
自知识蒸馏
嵌入模型被训练为将正样本与负样本区分开,对于每种检索方法,期望为查询的正样本分配比负样本更高的得分。训练过程旨在最小化 InfoNCE 损失,其一般形式为:
其中,和
分别表示查询q的正样本和负样本;s(⋅)是
中的任意一个函数。
不同检索方法的训练目标可能会相互冲突。为了促进多种检索功能的优化,在训练Bge-M3时,提出在自知识蒸馏的基础上统一训练过程。具体步骤如下:
1. 集成预测得分:将不同检索方法的预测得分整合为更准确的相关性得分:
2. 计算损失函数:计算的加权和作为没有自知识蒸馏的损失:
3. 修改损失函数:使用集成得分作为教师,修改每个检索方法的损失函数:
其中,p(⋅)是 softmax 激活函数;是
中的任意一个成员。
4. 最终损失函数:
- 进一步整合并归一化修改后的损失函数:
- 最终损失函数为:
训练过程分为两个阶段:
- 预训练阶段:使用大规模无监督数据对文本编码器(XLM-RoBERTa 模型,通过 RetroMAE 方法调整)进行预训练,仅以对比学习的基本形式训练密集检索。
- 微调阶段:应用自知识蒸馏,对嵌入模型进行微调,以建立三种检索功能。在此阶段使用标注数据和合成数据,并引入硬负样本。
高效批处理
嵌入模型需要从多样化和大规模的多语言数据中学习,以充分捕捉不同语言的通用语义。为了确保文本嵌入的区分性,模型需要尽可能保持大批量(引入大量批内负样本)。然而,由于 GPU 内存和计算能力的限制,传统的做法是将输入数据截断为短序列,以实现高训练吞吐量和大批量。这种方法虽然有效,但对于BGE-M3-Embedding 来说并不适用,因为它需要同时处理短序列和长序列数据,以有效支持不同粒度的输入。
为了克服上述挑战,本文提出了一种高效的批处理策略,具体包括以下几个关键步骤:
- 按序列长度分组:训练数据按序列长度分组,生成小批量时从同一组中采样,显著减少序列填充,提高了 GPU 的利用效率。
- 固定随机种子:在为不同 GPU 采样训练数据时,固定随机种子,确保负载平衡并减少每个训练步骤的等待时间。
- 子批量处理:在处理长序列训练数据时,将小批量进一步划分为子批量,然后使用梯度检查点(gradient checkpointing)技术,迭代编码每个子批量,显著减少内存占用的同时显著增加了批量大小。
- 跨 GPU 广播:在分布式训练环境中,将来自不同 GPU 的嵌入结果进行广播,使每个 GPU 都可以获取所有嵌入结果。通过广播,每个 GPU 上的模型可以访问更多的负样本,从而显著扩大了批内负样本的规模,有助于提高嵌入的区分性,进一步提升模型的性能。
2、实验结果
多语言检索
使用 MIRACL,包含 18 种语言的临时检索任务,每个任务由同一语言的查询和段落组成。
- Dense 方法:M3-Embedding 在密集检索上表现出色,平均性能优于所有基线方法。例如,在英语上与 E5mistral-7b 相比,M3-Embedding 在其他语言上表现更好。
- Sparse 方法:M3-Embedding 的稀疏检索功能优于传统的 BM25 方法。
- Multi-vec 方法:多向量检索进一步提升了性能,通过细粒度的交互计算相关性得分。
- 混合方法:Dense+Sparse 和 All 方法通过结合不同方法的优势,进一步提升了检索性能,其中 All 方法表现最佳。
跨语言检索
使用 MKQA,包含 25 种非英语语言的查询,目标是从英文维基百科中检索包含答案的段落。
M3-Embedding 在跨语言检索任务中表现出色,尤其是在密集检索和多向量检索方面。结合不同方法的混合检索进一步提升了性能。此外,M3-Embedding 在低资源语言上的表现相对稳定。
多语言长文档检索
BGE-M3可以支持长达8192的输入文档,同时长文档检索的效果要显著优于此前的基线模型。非常有意思的是,从实验结果可以观察到,稀疏检索(Sparse)的效果要显著高于稠密检索(Dense),这说明关键词信息对于长文档检索极为重要。
消融实验
自知识蒸馏和多阶段训练对 M3-Embedding 的性能提升起到了关键作用。通过这些技术,M3-Embedding 在多语言检索、跨语言检索和长文档检索任务中均表现出色。
3、总结
BGE-M3-Embedding模型的目的是构建一个多功能、多语言、多粒度的文本嵌入模型,能够高效地支持各种信息检索和自然语言处理任务。它的意义在于:
- 多功能性:统一了多种检索功能,适应不同的检索需求。
- 多语言支持:支持超过100种语言,适用于全球化的信息检索和多语言对话系统。
- 多粒度处理:能够处理从短句子到长文档的不同粒度的输入,适应各种文本长度的需求。
- 实际应用:在搜索引擎、问答系统、推荐系统、文档分析和多语言对话系统等领域具有广泛的应用前景。
从评测结果来看,BGE-M3全面应用于检索任务,包括训练数据也全是检索的相关任务,非常适合检索增强任务。
在实际应用中,BGE-M3也是text-embedding模型中下载量最高的模型,是目前最适合检索增强任务的开源文本嵌入模型之一。