2023-arXiv-GTE: 基于多阶段对比学习的通用文本嵌入模型

arXiv | https://arxiv.org/abs/2308.03281

gte-Qwen2-1.5B-instruct | https://www.modelscope.cn/models/iic/gte_Qwen2-1.5B-instruct

gte-Qwen2-7B-instruct | https://www.modelscope.cn/models/iic/gte_Qwen2-7B-instruct

摘要：

我们提出了GTE，一种通过多阶段对比学习训练得到的通用文本嵌入模型。随着近期将各类自然语言处理任务统一为单一格式的进展，我们通过在多源数据集上进行对比学习，训练了一个统一的文本嵌入模型。通过在无监督预训练和有监督微调阶段显著增加训练数据量，我们的模型在性能上大幅超越了现有的嵌入模型。值得注意的是，即使参数规模相对较小（110M），GTE_base 在大型文本嵌入基准测试中不仅超越了 OpenAI 提供的黑箱嵌入API，还超越了参数规模大 10 倍的文本嵌入模型。此外，无需针对每种编程语言单独进行额外微调，我们的模型通过将代码视为文本，在代码检索任务上超越了以往同规模的最佳代码检索器。

一、引言

文本嵌入已成为众多自然语言处理任务中不可或缺的组成部分，例如文本分类、文本检索、问答系统及对话系统。这些嵌入模型通过低维向量表示文本，并通过向量运算捕捉其相似性。近期大型语言模型（LLMs）的出现引发了基于文本嵌入模型的检索增强系统的广泛关注。

在文本表示研究领域，以往的文本嵌入模型主要聚焦于特定任务，其针对单一任务定制的训练策略或模型在其他情境下可能表现欠佳。近期，研究焦点转向利用大量未标注的网络数据通过无监督对比预训练，结合任务特定数据、提示或指令，开发更为全面的文本表示模型，以缓解微调过程中的任务冲突。然而，现有研究的一大局限在于依赖内部数据进行预训练，这在使用预训练模型权重或API时形成了瓶颈。

本研究提出了一种简洁的方法，仅通过对比学习在开源数据上构建**通用文本嵌入（GTE）**模型。

首先从多种数据源中收集大规模的无监督文本对数据集，用于对比预训练。
其次，为了进一步提升学习到的文本表示质量，我们从多个来源获取了带有高质量人工标注的文本对，用于对比微调。

二、模型方法

2.1 模型基本框架

我们嵌入模型的核心是一个深度 Transformer 编码器，可以使用预训练的语言模型如 BERT 进行初始化。模型遵循经典的双编码器架构，在语言模型生成的上下文 token 之上采用均值池化操作。

给定一段由 n 个 token 组成的文本 $x=(x_1,…,x_n)$ ，嵌入模型 $E$ 将该文本转换为一个低维密集向量 $x=E(x)\in R^d$ 。

首先采用一个语言模型来获取深度上下文化的 token 表示：

$\pmb{\text{h}}=\text{LM}(x)\in R^{n\times d}$

随后在第一维度上应用轻量级的均值池化操作，以获取文本表示：

$\pmb{\text{x}}=\frac{1}{n}\sum_{i=1}^n\pmb{\text{h}}_i\in R^d$

文本表示通过对比目标进行学习，以区分语义相关的文本对与不相关的文本对，正样本对和负样本对 $q,d^+, d^−)$ 。对于一个查询 $q$ ，一个相关文档 $d^+$ ，以及一组不相关文档 $D_-=\{d^−_1,…,d^−_n\}$ ，常见对比目标是 $\text{InfoNCE}$ 损失。

$L_{\text{cl}}=-\log\frac{e^{s(q,d^+)/\tau}}{e^{s(q,d^+)/\tau}+\sum_{i=1}^n e^{s(q,d^+)/\tau}}$

其中， $s (q, d)$ 通过向量距离估计两段文本 $q$ 和 $d$ 之间的相似度， $\tau$ 为温度参数

为了获取能够在广泛场景中应用的高质量文本嵌入，我们从多种格式和领域中汇编了一个广泛的文本对数据集。随后，采用改进的对比损失方法，以多阶段的方式对该数据集进行训练。

2.2 无监督预训练数据

弱监督的文本相关性数据在公开可访问的网络资源中易于获取，例如问答论坛上问题与答案之间的内在联系。这些数据无需人工标注即可广泛收集，从而有效辅助文本表示模型的训练。

为确保嵌入模型的通用性，我们探索了多种文本对提取资源，包括网页（CommonCrawl、ClueWeb）、科学论文（arXiv、SemanticScholar）、社区问答论坛（StackExchange）、社交媒体（Reddit）、知识库（Wikipedia、DBPedia）以及代码仓库（StackOverflow、GitHub）。此外，我们利用某些数据集中存在的超链接来促进文本对的提取。

2.3 有监督微调数据

在有监督的微调阶段，我们采用规模相对较小的数据集，这些数据集包含人工标注的两段文本之间的相关性，以及通过额外检索器挖掘的可选硬负例，以构成文本三元组。

为了同时处理对称任务（语义文本相似性）和非对称任务（段落检索），我们从多种任务和领域收集数据，包括网络搜索（MS MARCO）、开放域问答（NQ）、自然语言推理（SNLI）、事实核查（FEVER）和释义（Quora）。

2.4 训练细节

2.4.1 数据采样

在无监督预训练的初始阶段，不同数据源在训练实例数量上往往存在显著差异。为解决这一不平衡问题，我们采用多项式分布从不同数据源中采样数据批次，同时考虑其各自规模。

假设整个预训练数据集 $D$ 由 $m$ 个不同子集 ${D_1,…,D_m\}$ 组成，每个子集的大小记为 $n_i=|D_i|$ ，在每次训练迭代中，从第 $i$ 个子集 $D_i$ 中采样数据的概率为：
$p_i=\frac{n^\alpha_i}{\sum_{j=1}^m n^\alpha_j},\alpha=0.5$
此外，为防止模型仅学习特定任务的捷径进行判别，我们确保一个批次内的所有训练实例均来自同一任务。

2.4.2 改进的对比损失函数

在使用对比目标时，研究者通常复用批次内的文档作为负样本候选，以提高训练效率。本文采用了一种改进的对比学习目标，该目标具有双向性，并通过同时利用批次内的查询和文档来扩大负样本规模。

考虑一批正文本对样本 $B=\{(q_1,d_1),(q_2,d_2),…,(q_n,d_n)\}$ ，改进的对比损失函数：
$L_{\text{icl}}=-\frac{1}{n}\sum_{i=1}^n\log\frac{e^{s(q_i,d_i)/\tau}}{Z}\\ Z=\sum_j e^{s(q_i,d_j)/\tau}+\sum_{j\not=i} e^{s(q_i,q_j)/\tau}+\sum_j e^{s(q_j,d_i)/\tau}+\sum_{j\not=i} e^{s(d_j,d_i)/\tau}\\$
配分函数 $Z$ 的前两项用于查询到文档的对比，而后两项则用于反向对比。

$s (q, d)$ 采用余弦相似度作为距离度量，温度参数 $\tau=0.01$ 。

2.4.3 训练和评估

**第一阶段：仅使用批次内负样本进行对比预训练。**通过包含更多负样本来减少训练与推理之间的差距，并更好地逼近基础学习目标。

最大序列长度限制为128，在所有GPU上分配负样本的使用。
联合采用了fp16自动混合精度训练、deepspeed ZeRO第一阶段以及梯度检查点等流行技术，以降低内存成本并将批次规模扩展至超过一万。
进行了 50,000 步的预训练，大致相当于在整个预训练数据上完成一个周期。
仅调整了学习率以确保较大模型的收敛性，采用了 AdamW 优化器，并在训练初期 5% 的步骤中实施线性学习率衰减和预热期。

小型、基础和大型三种不同规模的模型，分别以 MiniLM 的小型模型和 BERT 的基础和大型模型进行初始化。

**第二阶段：利用监督数据和困难负样本进行对比微调。**由于困难负样本已能提供学习目标的可靠梯度估计，因此无需采用大批次规模。

全局批次大小为128，训练组大小为16，其中包含一个正例样本，其余则为困难负样本或随机负样本。
最大序列长度增加至512，以更好地处理较长文本。
在微调过程中，学习率降低了十倍。
模型在收集的数据集上进行了单轮微调。
采用了改进的对比损失函数，将批次内文本也纳入负样本候选。

三、实验结果

3.1 零样本文本分类

将文本分类重新定义为基于嵌入的相似性匹配问题，输入文本直接转换为嵌入，标签被表述为相应的文本以获得标签嵌入。输入嵌入与标签嵌入之间的距离通过它们的内积来衡量，与输入文本嵌入距离最近的标签被视为分类结果。

SST-2 二元情感分类任务，两种类型的标签表述器：

基础版本：使用情感词**“积极”或“消极”**来表示相应的标签。
提示版本：使用模糊提示模板，如**“这是一个积极/消极电影评论的例子”**。

3.2 无监督文本检索

文本检索需要从大规模候选集中检索出最相关的文档。采用 BEIR 作为零样本无监督文本检索的评估基准，BEIR 是一个异构信息检索基准，包含不同格式和来自不同领域的检索任务，使用公开可用的15个数据集进行评估。

3.3 大规模文本嵌入基准

大规模文本嵌入基准（MTEB）是一个综合性的半监督基准，整合了有限的监督数据用于评估。

本文评估了英语子集，涵盖了七个不同任务下的56个英语数据集，包括文本分类（Class.）、文本聚类（Clust.）、成对分类（Pair.）、文本重排序（Rerank.）、文本检索（Retr.）、语义文本相似性（STS）以及摘要生成（Summ.）。MTEB采用的评估指标分别为准确率、V-measure、平均精度、平均精度均值（MAP）、归一化折损累计增益@10（nDCG@10）以及斯皮尔曼相关系数。