MGTE系列模型，RAG实施中的重要模型-CSDN博客

检索增强生成（Retrieval-Augmented Generation, 简称RAG）正日益成为大模型应用中的热门范式。这种创新方法将检索与生成相结合，使得大模型在提供回答时，能够充分借助外部知识库的支撑，从而获得更为准确且丰富的内容。这不仅有效地减轻了大模型常见的误解和数据隐私问题，也提升了实时响应的能力。

在RAG实施过程中，文本表示模型（Embedding model）和排序模型（Reranker model）是至关重要的两大模块。两者的目标一致，皆为检索与用户问题相关的文档，但在实现方法上却有所不同：文本表示模型首先计算每段文本的向量表示，然后通过余弦距离等方法计算它们之间的相关性分数。因此，所有文档的向量表示可以在离线阶段提前计算，在线上只需处理用户查询的向量，通过高效的向量搜索引擎迅速获取相关候选文档。而排序模型则直接将文本对作为输入，通过更精细的计算方法输出它们的相关性分数，实现更准确的排序效果。需要注意的是，排序模型的计算复杂度较高，因此主要适用于小规模的候选集合。

随着RAG技术的不断发展，系统的功能也在日益增强。早期的系统主要关注向量召回，而如今更复杂的模块被逐步引入，进一步提升了RAG系统的整体效果。同时，随着应用场景的拓展以及大模型处理上下文长度的能力增强，对多语言检索、跨语言检索以及长文本处理的需求也愈加迫切。

此前，通义实验室推出了GTE（General Text Embedding）系列文本向量模型，涵盖了基于BERT架构的模型及基于Qwen LLM系列训练的LLM embedding模型，如gte-Qwen2-1.5B-instruct与gte-Qwen2-7B-instruct。目前，基于双向注意力的Encoder-only结构的模型在同一规模下相较于Decoder-only模型在召回和排序效果上明显更优。然而，当前基于Encoder-only的Embedding和Ranking模型仍面临一些来自BERT时代的遗留问题，例如最长上下文长度仅为512，以及预训练阶段使用的语料明显不足。为解决这些问题，GTE模型开发团队从零开始训练了一种能支持长上下文和多语言的Encoder-only基础模型，并在此基础上推出了最新版本的GTE-MultiLingual系列模型（简称mGTE模型）。该系列模型具备以下显著特点：

高性能：在多个数据集上与同规模开源模型的对比中，对比效果领先。
长文档支持：Embedding和Reranker均可处理8k token文本长度，且支持通过ntk-rope等方法扩展到更长的上下文。
多语言支持：模型支持75种语言，涵盖当前主要大模型所支持的所有语种。
弹性向量表示（Elastic Embedding）：模型支持输出128-768维度之间的任意向量表示，以便在性能和存储成本之间取得最佳平衡。在128维的情况下，与768维相比，召回性能损失小于2%，同时节省6倍的存储空间。
稀疏向量表示（Sparse Embedding）：模型可以输出句子中每个单词的词权重作为稀疏表示，适用于需要精确匹配的场景。

图1 文本表征和文本排序模型架构示意图

2. 模型构建

mGTE系列模型构建流程如图2所示，首先，训练了支持长下文的多语言Encoder-only底座模型GTE-base-multilinguish。并在底座基础上继续训练文本表示模型gte-multilingual-base和排序模型gte-multilingual-base。

图2 模型训练过程示意图

3. 底座预训练

3.1 模型结构

为了提升模型多语言以及长文本领域相关的能力，该系列模型参考了目前Decode-Only架构大语言模型训练过程中一些常见的技巧，对原始的BERT架构做了以下几点改动，具体模型结构如图3所示

位置编码: 将BERT模型中采用的绝对位置embedding方式改为了旋转位置编码RoPE [1]，以便能更好的支持长上下文的训练，同时保持上下文长度扩展的可能性。
激活函数: 将BERT模型中线性层（FFN）部分改为了GLU（gated linear unit，[2]），这也是在LLM训练过程中已经经过充分验证能有效提升模型训练稳定性的技巧。

此外，为了满足多语言和长文本处理能力的需求，模型使用了XLM-Roberta[3]系列的词表。

图3 GTE底座模型结构示意图

3.2 训练数据

底座模型的训练数据来网络上公开的多多言语料数据，包括C4，Skypile，mC4，CuleturaX，Wikipedia，books等。训练过程中，通过对数据进行过滤，清洗，采样等操作，最终得到了1028B的Token（基于XLM-R tokenizer）作为最终的训练数据。其中主要语言的比例如下图所示（占比小于1%的合并为 others），各语言的数据量可以参考mGTE论文[16]中的附录。

3.3 训练过程

与传统的Encoder-only模型训练方法一样，预训练阶段mGTE使用MLM（Masked Language Prediction）作为预训练阶段的loss。在此基础上，该模型做出了以下优化提升模型训练的效率：

数据采样：训练过程中，为保证每个Batch内的数据均来自同一种语言，每次训练过程会从所有语言中根据概率采样某种语言的数据，具体每一种语言的采样的概率

计算方法如下：

其中n代表的是该语言的数据量，

。

多阶段预训练：为提高训练效率，首先将数据截断为2k长度进行训练，之后在此基础上将数据截断到8k进行继续训练，并将RoPE的base参数由10000设置为160,000。
Unpaddding Model：训练过程中，由于并不是所有的文档都都是统一的长度，mGTE利用了unpadding技术来避免在计算过程中对padding部分的token进行计算，通过这一技术可以有效的提升训练效率，目前常见的flash-attention与xformers库中均支持该操作。

预训练阶段，模型均使用BF16精度进行训练。

3.4 底座模型效果评测

为了验证预训练底座模型的效果，mGTE-mlm在多语言评测榜单XTREAM[4]（共包含50种不同的语言）和英文GLUE[5]数据上对比了2K阶段与8K阶段训练模型与之前Encoder-only架构的同尺寸多语言模型的效果，可以看到在绝大多数任务上mGTE-mlm都相较于过去的模型有更好的效果。

表1 XTREME-R多语言数据评估

表2 英文GLUE数据集评估

4. Embedding模型

目前基于Encoder-only架构的文本表征模型基本上都采用双阶段训练的范式(弱监督训练与监督训练)[6]，以提升模型的泛化性和效果。

4.1 弱监督训练

在弱监督训练阶段，mGTE-embedding利用从网络上爬取的大规模文本对数据集，例如网页的标题与正文、QA论坛中的问题与答案等，这些数据无需人工标注，相对容易获取，但质量可能较低。这一阶段通过对比学习方法进行训练，使模型具备基本的嵌入表示能力。经过数据清理和整理，共收集约28亿多语言弱监督数据。在该阶段中，mGTE使用了主流的对比学习方法进行模型训练。

4.2 监督训练

监督训练阶段，mGTE使用了经过人工标注的高质量文本对数据对模型进行进一步训练，以提升模型的文本表示能力。其中，中文监督数据包含Dureader，Simcluse，Multi-CPR等在内的6个数据集，共计200万标注数据，英文监督数据集使用了包括Msmarco，Nq，NLI等在内的7个数据集，共计140万数据。此外还包括MLDR，MIRACL，Mr.TyDi三个多语言数据集：共计12万数据，具体的数据集情况请参考论文附录部分。

在监督训练阶段，除了基础的连续向量对比学习损失，mGTE引入了两个额外的表征特性:

弹性维度表示：弹性维度表示模型可以输出不同维度的向量表示，以平衡在索引存储和模型效果之间的Trade Off。目前许多方法开源与闭源模型均通过MRL学习[7]支持该特性。模型可以输出不同维度的向量表示，从而在索引存储和模型效果之间找到平衡。在模型训练过程中，令D代表一个整数列表，并分别取模型最终向量表示的前k维度（k属于D）进行标准化并计算对比学习损失。最终，这部分的损失是各个维度对比损失的平均值。
Sparse向量表示：与连续向量(Dense)表示方法不同，Sparse表示是指通过深度模型计算文本中每个单词的词权重，通过两段文本之间相互匹配的词的权重乘积的和作为两段文本的相似度，可以认为是传统BM25等检索方式的一种扩展。相比于Dense表示，Sparse表示在需要精准匹配（如型号名，品牌名，年份匹配等）场景，以及长上下文检索场景往往有更强的性能。训练过程中，mGTE在每个位置token最后一层的输出增加了一层线性层（使用Relu作为激活函数）得到文本的Sparse表示，并使用对比学习损失作为Sparse向量表示的损失函数。

因此, 模型最终的训练损失函数为MRL部分损失函数与Sparse部分损失函数的加权和：

在训练过程中，该模型采用使用了两个额外的策略来保障模型性能与训练效率：

Hard Negatives构造：之前的研究工作表明对比学习的效果高度依赖于负样本的数量和质量[8]。基于此，mGTE首先利用在弱监督阶段训练得到的模型，从每个查询（query）中挖掘出难负样本（排序靠前的负样本）。此外，mGTE还将同一个批次（batch）内其他查询的文档（包括相关文档和不相关文档）用作负样本，从而扩大负样本的数量和多样性。为了进一步保证负样本的质量，以确保每个Batch内的数据都来自同一数据集。
动态Batch大小：为了提高长上下文的检索能力并最大化训练效率，mGTE采用了动态 Batch 大小的策略。首先，根据每个文档的长度对训练数据进行分组，并针对不同长度的数据使用不同的batch size。在训练过程中，首先将一个Batch切分为多个子Batch（sub-batch）。通过Gradient-checkpoint方法，可以依次计算每个子Batch的Embedding，并通过合并所有子Batch的表示得到完整Batch的Embedding。该方法可以使Batch size增加2到5倍，显著提升训练效率。

Embedding与Ranking的训练过程都是基于FP16精度进行的，使用了DeepSpeed zero0与Gradient checkpoint来节省内存。

4.3 效果评测

检索效果评测

为了评测文本表征模型的检索效果，尤其是模型的多语言和长文本处理能力，mGTE模型主要在以下几个数据集进行评测

MLDR[9]: 多语言长文档检索评测集, 包括13个语种数据。
MIRACL[10]: 多语言检索评测集合,包含18个语种数据。
MKQA[11]: 跨语言检索评测集，包含25个不同的语种。
BEIR[12]: 英文多领域检索评测集合。
LoCo[13]：英文长文档检索评测集合。

表3展示了在这5个数据集上mGTE模型和同规模模型的效果对比:

受益于原生的长文本底座训练，mGTE表征模型在长文本检索效果上明显优于其它模型。
在短文本检索场景, mGTE对比同规模的模型效果大幅度领先，对比更大规模的模型效果也很接近。
Sparse向量检索效果在大部分场景优于BM25，特别在长文档场景对比现有Dense向量检索有明显效果优势。

表3 检索数据集效果对比

多任务文本表征效果评测

MTEB[14]是一个涵盖多任务文本表示的通用评测数据集，英语、法语和波兰语这四种不同语言上对mGTE模型与其他模型的性能进行了对比。类似于检索任务，mGTE模型在与开源社区中同规模的Encoder-only系列模型进行比较时表现出非常好的效果。当然，与更大型的基于LLM的模型相比，mGTE仍存在明显差距。然而，考虑到mGTE小模型在推理性能方面的优势，其在实际应用场景中应具备更大的发挥空间。

表4 MTEB多任务多语言效果效果对比

弹性向量表示

通过进行弹性向量表示，可以大幅提升文本处理与信息检索的效率和效果，使高维文本数据的处理更加可行和高效。在MTEB英文评测中，mGTE对比了mGTE模型在不同维度表征下的表现，结果与OpenAI的弹性表示模型类似。降低同一模型的向量表征维度虽然会导致部分效果的损失，但在512维以上的情况效果损失仍在可接受范围内。

图4 MTEB多任务多语言效果效果对比

5. 排序模型

在训练排序模型时，mGTE-reranker继续采用对比学习损失函数。论文中介绍通过实现发现，弱监督训练阶段对排序模型效果的提升非常有限。因此，在最终的排序模型训练过程中，mGTE-rerank仅使用监督数据进行微调。由于Ranking模型需要以文本对作为输入来计算相关性分数。此外，排序模型的超参数设置与文本表示模型保持一致。

5.1 效果评测

类似地，mGTE-reranker对MLDR、MIRACL、MKQA和BEIR数据集上的排序结果进行了评估。所有的排序模型都基于mGTE-TRM-base模型生成的向量召回的结果top100进行重新排序。具体结果如下：

图5 排序模型结果对比图

图5展示了不同排序模型的检索效果：

各个排序模型相比召回模型均表现出更佳的性能，证明了在检索链路中引入排序模型的必要性。
与同尺寸甚至更大尺寸的模型相比，mGTE-reranker-base模型在各个数据集上均取得了相当甚至更好的效果，尤其是在多语言长文档的检索场景中。

6. 模型使用

模型的使用方法可以参考Huggingface上的样例：

6.1 Embedding模型：

# Requires transformers>=4.36.0

import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer

input_texts = [
    "what is the capital of China?",
    "how to implement quick sort in python?",
    "北京",
    "快排算法介绍"
]

model_path = 'Alibaba-NLP/gte-multilingual-base'
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModel.from_pretrained(model_path, trust_remote_code=True)

# Tokenize the input texts
batch_dict = tokenizer(input_texts, max_length=8192, padding=True, truncation=True, return_tensors='pt')

outputs = model(**batch_dict)

dimension=768 # The output dimension of the output embedding, should be in [128, 768]
embeddings = outputs.last_hidden_state[:, 0][:dimension]

embeddings = F.normalize(embeddings, p=2, dim=1)
scores = (embeddings[:1] @ embeddings[1:].T) * 100
print(scores.tolist())

6.2 Ranking模型：

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained('Alibaba-NLP/gte-multilingual-reranker-base')
model = AutoModelForSequenceClassification.from_pretrained('Alibaba-NLP/gte-multilingual-reranker-base', trust_remote_code=True)
model.eval()

pairs = [["中国的首都在哪儿"，"北京"], ["what is the capital of China?", "北京"], ["how to implement quick sort in python?","Introduction of quick sort"]]
with torch.no_grad():
    inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512)
    scores = model(**inputs, return_dict=True).logits.view(-1, ).float()
    print(scores)