大模型时代如何得到更好的embedding表征？

导读

世界上本没有「不可能」，存在的「不可能」只是自己视角下的「不可能」。——爱工作的小小酥

LLM如何学习更好的embedding？

以往的训练embedding的模型都是以BERT架构为基础，使用MLM损失训练，得到模型可以在各种下游任务上微调，例如分类。

在大模型时代，我们经常采用语言模型损失建模，以预测下一个token的方式进行训练，这样训练的模型在很多任务上都超过了以前的模型，但这样的模型怎么提取embedding特征，用于其他检索相关任务呢？

最近注意到MTEB排行榜上多了很多基于现有的大模型训练的模型，在此整理一下embedding技术的发展，以备不时之需，但读的时候发现并不是所有论文都使用了LLM。

阅读了排行榜前排的工作，发现现在embedding模型主要有以下特点：

1、使用INfoNCE loss。

在之前训练BERT的时代，如果进行预训练，一般是用MLM任务，可能最近大家评测的很多任务或者大模型RAG的兴起，导致检索任务变得额外重要。下面的损失函数基本是用INfoNCE，个别使用了cosent loss（比较类似INfoNCE，是一个可以用于排序的loss）。

2、多任务训练。

使用INfoNCE loss进行训练的话，需要构建（query，正样本，负样本）三元组。对于检索类的任务，当然很容易构造，而对于分类、聚类的任务，就要想各种办法构造出这种格式，在下面的论文中主要涉及两种方式：

• 将分类、聚类的类别标签作为数据匹配的正样本或者负样本。（这个相对较多）
• 将分类、聚类的同类别其他数据作为正样本，不同类别的其他数据作为负样本。

其次，在多任务训练的时候，经常针对不同任务采取不同的损失函数，或者尽管都采用INfoNCE，对于非检索任务，一般不再采用in batch negative。

3、Matryoshka Representation Learning

俄罗斯套娃表示训练方法，其因为可以同时训练多个维度的向量表示，又不会影响性能的优势，在目前的embedding模型中经常被采用。（而且很多都是最大到1972）

4、多阶段训练

可能借鉴于目前LLM的训练流程，目前很多不使用LLM的embedding模型，也采用了多阶段进行训练，一般是划分为2阶段，先使用低质量文本对训练，然后使用人工标注的检索类数据集训练。

5、困难负样本挖掘

困难负样本挖掘一直是对比学习的重点，在下面的几篇论文中，很多也很强调困难负样本的重要性，但主要方法也比较受限，主要是以下几点：

• 单独训练一个embedding模型，给负样本打分，选择高分数据作为困难负样本（用的较多）
• 动态在训练过程中更换困难负样本，计算方式类似上面

6、使用合成数据

在大模型时代，经常使用LLM造数据，在embedding任务中，大家也是进行各种prompt工程构建检索类数据，主要是分为两阶段构造，先让模型「头脑风暴产生主题」，然后「根据主题生成三元组数据」。

7、文本表示方法

在这个方面和之前差不多，主要有以下几种：

• 使用最后一层特征的 mean pooling
• 使用一个特殊的token。[CLS]或者[EOS]
• 后面接一个attention pooling 层。

Piccolo2

Piccolo2: General Text Embedding with Multi-task Hybrid Loss Training

https://arxiv.org/abs/2405.06932

概要

这篇论文核心点主要有：

• 将不同的任务使用「不同的损失函数」进行训练，为检索、排序、分类、聚类配置不同的loss，从而兼顾每一种任务的特点，在各个任务上都可以达到最优的效果。
• 在训练过程中加入Matryoshka Representation Learning，从而训练出动态维度的向量。其中最高维度到1972。
• 使用合成数据训练。
• 负样本挖掘。使用piccolo-base-zh使用相似度排名在50-100的随机15个作为困难负样本。

模型结构依旧是bert，训练完的参数量为300M。

训练方式

（1）Retrieval and Reranking Loss

对于检索和排序任务，采用标准的INfoNCE loss，并使用in batch negative。

（2）Semantic Textual Similarity（STS） and PairClassification Loss

在这种任务中，其label一般不是一个绝对的值，例如相似度分数，如果将这种任务单纯的表示为INfoNCE中的三元组的形式，就丢失了这部分信息，因此，对于这些任务，作者使用了cosent loss。具体原理可参考下面的链接：

CoSENT（一）：比Sentence-BERT更有效的句向量方案 – 科学空间|Scientific Spaces

https://kexue.fm/archives/8847

（3）Classification and Clustering Loss

在分类和聚类任务中，没有样本对的概念，而是一个标签。作者对于这种任务，是将其label和数据作为一个正样本pair，当前数据和其他label作为负样本pair，采样标准的INfoNCE loss进行训练，但不再使用in batch negative。

最终的损失函数为上述3个loss相加。

合成数据

为了构造出丰富多样的检索相关的数据，作者分为2个步骤进行，分别为生成话题、根据话题生成样本对。最后生成了200k的数据，下面为整个训练过程中采用的数据量。

Conan-embedding

Conan-embedding: General Text Embedding with More and Better Negative Samples

https://arxiv.org/abs/2408.15710

概述

困难负样本一直是对比学习训练过程中的一个关键点，困难负样本的好坏有时会非常影响模型的性能。现有的挖掘困难负样本的方法基本集中在训练之前，经常使用训练好的embedding模型计算负样本的相似度，将高相似度的负样本筛选出来作为困难负样本。但这样一次性决定困难负样本的方法是最优的吗？模型训练过程中会不会改变负样本的难度？

在这篇论文就是集中在这个问题上研究的，设计了一种在训练过程中挖掘负样本的方法，可以根据模型训练的状态选择合适的困难负样本进行学习。其次，为了进一步扩大负样本的数量，提出了Cross-GPU Batch Balance Loss (CBB)。

在模型结构上，依旧采用bert，使用的是BERT large模型，并且也使用了Matryoshka Representation Learning (MRL) 技术训练动态embedding，最大到1792。

训练方式

第一阶段，使用bge-large-zh-v1.5筛选出高质量数据（阈值高于0.4），这里的数据是常规文本对和合成数据，总共有0.75 billion。采用标准的INfoNCE loss 进行训练，并使用in batch negative。

第二阶段，主要使用检索和 STS (semantic textual similarity)任务数据，并采用不同的损失函数，其中检索使用标准的INfoNCE loss，STS使用CoSENT loss。

动态困难负样本

为了在训练过程中动态的选择困难负样本，作者在每100step进行检查一次，查看当前困难负样本乘以1.15是否小于原始的score并且绝对分数小于0.8，如果存在，则进行更换困难负样本。

对于第  轮，更换的时候，选择  到  之间的负样本，其中 为困难负样本的个数。（不懂为啥这样选？？）

Cross-GPU Batch Balance Loss

由于训练过程中针对不同任务采用了不同的损失，在以往的训练过程中，是每一个batch都采样同一来源的数据，但这样训练会导致不稳定性，具体可看下图，cross是合并后的损失。

因此，作者将两种任务合并训练，损失为两者的相加，使用不同的gpu来平衡batch的大小。

GTE

Towards General Text Embeddings with Multi-stage Contrastive Learning

https://arxiv.org/abs/2308.03281

概要

阿里GTE系列模型的第一个版本，和以往训练bert不一样的地方在于：

• 收集大量数据
• 分为两阶段进行训练
• 采用升级版对比学习损失，扩大了负样本
• 更关注检索任务

模型结构依旧采用bert，感觉主要是靠大数据量造就了一个比较好的效果。

模型结构

模型的整体结构没有作出改变，加载预训练的MiniLM、bert分别训练出3种大小的模型，分为small、base、large，不同规模的模型参数如下：

损失函数

均衡采样

在预训练阶段由于混合了多种来源的数据，而不同来源的数据数量存在较大差异，为了缓解这种多种数据源数据严重不均衡的问题，使用了基于原始数据量的采样。具体来说，在每一轮的训练过程中，从每个来源采样的概率为其在所有数据中所占的一个大致比例，具体公式如下：

对于上述公式中的  ，论文中取的是0.5。并且为了不让模型「偷懒」，确保每个batch内都是相同类别的数据（后续的mGTE也沿用了这种采样方法）。

升级版InfoNCE

对于（query, document) 这样的数据对来说，原始的InfoNCE是in batch negative，即每一个query和同batch内的其他document构成负样本。因此，在常规的对比学习任务中，一般是batch越大越好，这样样本见到的负样本会更多，会更有利于模型学习，但batch又收到机器显存的限制。

因此，在这篇论文中，在不增大batch的情况下，为了增大每个样本的负样本数量，作者在原始InfoNCE loss的基础上做了进一步的改进，即不再仅仅使用「单向的负样本」，而是使用了「双向的in batch query」和「document negative」。

操作过程也比较简单，可以想一想对于一个batch内的所有数据来说，一条（query, document) 是不是和其他所有的（query, document) 都可以形成负样本，并且不限于（query, document) 这种格式。

那么，我们就可以在原始（query, document) 的基础上扩展（query, query) 和（document, document) 也构成负样本。

具体可以看一下下面的公式，对于一个包含（query, document) 的batch数据：

改进版的InfoNCE的负样本构成如下：

1. 当前query和其他document（原始InfoNCE的负样本）
2. 当前query和其他query（新增）
3. 其他query和当前document（新增）
4. 其他document和当前document（新增）

训练方式

第一阶段 无监督预训练

整篇论文主打多种数据、大批量数据训练，此部分就将各种来源的text pair数据进行了混合，数据量总共为788M，具体比例如下：

第二阶段 监督微调

为了在下游检索任务上得到更好的效果，此部分使用更高质量的人类标注pair数据，并使用其他的检索器抽取困难负样本，包括web search (e.g., MS MARCO),open-domain QA (e.g., NQ), NLI (e.g., SNLI), fact verification (e.g., FEVER), paraphrases (e.g., Quora)，总共大约3M pairs。

mGTE

mGTE: Generalized Long-Context Text Representation and Reranking Models for Multilingual Text Retrieval

https://arxiv.org/abs/2407.19669

概要

从模型的名称mGTE就可以看出，这篇论文主要着重点在「多语言」上，除此之外，论文还着重强调在「长文本」上的优越性。

和以往的BERT系列模型不同，此模型支持8192长度（之前都是512长度），并采用较为复杂的多阶段训练流程，同样和GTE类似，使用了大量的数据，最后不仅训练出召回模型，同时也得到了一个优秀的reranker模型。

由于在此时这个阶段，对于如何学习更好的embedding问题？已经出现了各种各样的技术，而在这篇论文中我们可以看到的主要有可以同时训练出多个维度embedding的Matryoshka Embedding方法（俄罗斯套娃embedding方法）、增强训练效率的unpadding技术、以及扩展长文本的Rotary Position Embedding编码技术。

模型结构

在模型结构上依旧采用常规的BERT结构，在其中加入了很多技巧性的工作。整体而言，如下图所示，总共产生了5个模型。这5个模型由一个复杂的训练流程串联起来，最终使用的主要是检索模型（TRM）和精排模型（reranker）。

虽然训练流程较长，但差异性并不是很大，下面先进行简短的总结一下：

1. 2k Text Encoder 和 8k Text Encoder 采用相同的数据和训练方法，只是输入的最大长度及个别训练参数不一样。并且采用以往bert使用的MLM损失进行训练。
2. 1k Text Embedder 是向检索任务过度的一个模型，采用常规infoNCE进行训练。
3. 8k TRM和8k Reranker采用相同的数据和训练损失，只是加载的基础模型不同（从图中也可以看出），且输入方式不同（毕竟一个双塔一个单塔）。

Text Encoder&Text Embedder

从bert之后，已经出现了很多如何更好的学习embedding的新技术，在这篇论文中主要用到有

• 大模型经常采用的Rotary Position Embedding编码方法可以扩展到训练未见过的长度
• unpadding技术可以更有效的训练
• 为了使用FlashAttention，将bert中的FFN改为gated linear unit (GLU)。

各种技术的原理可以再去查找相关资料，本文不再赘述，下面给出2篇unpadding相关的论文。

A Bidirectional Encoder Optimized for Fast Pretraining

https://mosaicbert.github.io/

https://arxiv.org/pdf/2208.08124

TRM&Reranker

对于检索和排序模型而言，整体的模型结构和上述类似，只不过是输入的方式不同。

在检索模型中，query和document是分别输入，使用CLS token向量作为整体的特征表示。

在精排模型中，query和document是拼接输入，同样使用CLS token向量作为整体的特征表示。

同样的，在微调这个阶段，在上述预训练的基础上，也额外使用了很多embedding的各种技术，包括Matryoshka Embedding方法（俄罗斯套娃embedding方法）、可以增强长文本能力的Sparse Representation loss。

训练方法

Text Encoder

在此部分，主要是为了增强模型对文本的理解能力，以及扩展长文本能力。因此，采用了常规的MLM的损失函数，将mask比例设置为30%，并使用了课程学习的训练方法，将文本长度从2k扩展到8k。

由于这篇论文主打「多语言」，因此在整个训练语料上加入了各种各样语言以及来源的数据，在这一阶段中，总共包括75种语言数据，总token数为1028B。

除此之外，在训练过程中，采用了一个采样的小技巧。为了平衡采样不同语言的数据，以防数据过多的数据被训练更多次，首先计算出每种语言所占的比例，然后采用下述多项式分布的方式进行采样，并且确保一个batch内是同一种来源的数据。

在这一过程中的训练参数如下：

Text Embedder

通过上述MLM损失训练后的模型，可以对文本有较好的语义理解能力，但在检索任务中可能效果没有那么好。而对于目前针对检索任务训练的各种召回和排序模型而言，基本采用对比学习的方式。

因此，为了增强模型在下游检索任务上的效果，作者同样采用对比学习损失做了继续预训练（CPT），损失函数为InfoNCE loss和in batch negative。

训练数据同样包括了各种语言各种来源的数据，主要有英文数据对、中文数据对、多语言、交叉语言指令和翻译数据对，总共2,938.8M pairs。

TRM&Reranker

得到一个适合检索的预训练模型之后，已经可以很好的做检索相关的任务了，但还不是一个reranker模型，无法做到精确排序。因此，在上述模型的基础上，作者进一步筛选高质量的pair数据，使用更贴合这个场景的训练损失，以得到更优的检索和排序模型。

这里的TRM和reranker训练过程大体一致，只是采用的预训练模型不一样，TRM使用1k Text Embedder作为基础模型，reranker使用8k Text Encoder 作为基础模型（官方也尝试使用1k Text Embedder，但无增益）。两者采用的训练数据如下：

真个微调过程中，在标准的infoNCE 损失的基础上，使用Matryoshka Embedding方法（俄罗斯套娃embedding方法）得到动态的embedding维度，使用可Sparse Representation loss进一步增强长文本能力。

（1）Matryoshka Embedding方法

原始的embedding模型中，输出是一个固定维度，例如768，最后使用768这个向量计算一个loss。在推理的时候，同样只能使用768维度的向量，无法压缩到512或者扩展到更高维度（除了重新训练模型）。

Matryoshka Embedding方法即是为了解决这种问题而提出的，这个直译「俄罗斯套娃embedding方法」。意思是将原先的768维度缩小为多个维度，例如[32, 64,128, 256,768]，训练的时候每一个维度后面单独使用一个线性分类器进行转换一下，然后分别计算损失，最后所有损失相加得到最终的损失。

Matryoshka Representation Learning

https://arxiv.org/abs/2205.13147

（2）Sparse Representation方法

这个方法来源于BGE M3-Embedding: Multi-Lingual, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation(https://arxiv.org/abs/2402.03216)。

在标准的INfoNCE loss的计算过程中，我们需要计算一个query和document的相似度，这里一般采用余弦相似度或者内积，在这个方法中，主要是对这个相似度计算方法进行改变的，具体来说，为每个token计算一个权重

然后计算「每个pair中相同的token的权重乘积和」作为当前pair的相似度，替换掉标准INfoNCE中的相似度计算方式。整个微调过程中的损失为：

其中，前一部分为上述的Sparse Representation loss，第二部分为Matryoshka Embedding 的loss。

NV-Embed

NV-Embed: Improved Techniques for Training LLMs as Generalist Embedding Models

https://arxiv.org/abs/2405.17428

概述

LLM学习embedding表示的方法都有什么？

目前很多论文基本将最后一个token的embedding作为整个文本embedding表示，但这种方法由于太依靠最后一个token的表示，容易带来偏差。

因此，在这篇论文中，提出一个Latent Attention Layer对LLM最后的embedding进一步编码，以得到信息融合程度更高的embedding表示。

那模型的其他设置有什么特殊的地方吗？

基本没有了，在训练方式上，和常规模型比较类似，采用2阶段训练方式。loss为和标准的INfoNCE，模型使用Mistral-7B，并使用lora进行训练，设置rank=16，alpha=32。

并且作者强调自己没有使用合成数据。

样本构造方式

对于（query，document）数据对，只在query前面加上指令，document前面不添加。后续也有工作指出推理的时候都可以添加，不会导致不一致。

训练方式

（1）第一阶段

采用公开检索数据，并使用in batch negative进行训练，训练之前使用单独训练的embedding 模型筛选高质量数据。

（2）第二阶段

这一部分采用非检索数据集，主要包括分类和聚类数据集，从MTEB英文训练集中筛选的各种分类数据集。为了将这种数据转为（query，document）数据对，作者按照标签进行归类，将同类别的数据作为正样本，不同类别的样本作为负样本（其他工作是将label作为匹配的数据）。

Improving Text Embeddings with Large Language Models

Improving Text Embeddings with Large Language Models

https://arxiv.org/abs/2401.00368

概要

现有的训练embedding表征的模型常常采用大批量的数据，2阶段方式进行训练。这篇论文不再使用那么复杂的流程，只使用合成数据，并且只有一个阶段训练流程，就达到了有竞争力的效果。

模型配置

和以往的LLM embedding模型差不多，同样使用Mistral-7b进行lora训练，并将lora设置为16。

损失函数依旧采用标注的INfoNCE loss和in batch negative。

最后训练的时候加入一些带标注数据，总共1.8M。

合成数据构造方法

为了构造出具有较强多样性的数据，作者将任务归为2类进行设计prompt，分别为非对称任务和对称任务。但都是使用gpt3.5-Turbo和gpt4进行生成。

（1）非对称任务

非对称任务主要包括短-长匹配、短-短匹配、长-短匹配、长-长匹配。对于这些子任务，分别设计prompt。并将生成数据的流程拆分为2步，先生成task名称，再针对task生成相应的数据。

（2）对称任务

作者将这部分任务分为monolingual semantic textual similarity (STS) 和 bitext retrieval。对于这种任务，直接使用1步生成相应的数据，具体可以看下面的示例。

为了进一步增强数据的多样性，在prompt构造的过程中，随机选择生成的长度和生成的语言，最终构造出500k数据。

GRIT

Generative Representational Instruction Tuning

https://arxiv.org/abs/2402.09906

概要

一个模型可以兼容生成和embedding任务吗？

常规大模型基本是以生成方式进行建模，如果要借用大模型架构构建embedding模型，一般是将生成损失改为INfoNCE，但这样建模的模型就失去了生成能力。

这篇论文就从这个角度出发，不丢弃任何一方，同时使用INfoNCE和生成损失训练模型，让大模型兼顾生成能力和embedding能力。

模型配置

但混合两种任务一起训练会有一个问题，对于embedding任务来说，一般都是采用双向注意力机制，让模型充分学习token之间的语义关系；但对于生成任务来说，一般采用 causal attention，后面的token只能看到前面的token。因此，作者将不同任务的attention矩阵进行了更改，并且embedding任务取最后一层特征的mean pooling作为最后的特征向量，具体的模型结构如下：

使用的基础模型和其他的模型差不多，也是Mistral 7B，同时也训练出了Mixtral 8x7B，都是全量训练（前面有的是lora训练）。

损失函数

对于embedding任务，采用标准的INfoNCE loss，并使用in batch negative。对于生成任务，采用保准的生成损失（预测下一个token），两者加权得到最终的损失。

Reranker

虽然模型没有单独为排序任务训练，但也进行了reranker的评测，借鉴的是另一篇论文：arXiv reCAPTCHA(https://arxiv.org/abs/2304.09542)，将召回的结果，拼接到prompt中，借助模型的生成能力进行排序，具体使用的prompt格式如下：

总结

随着大型语言模型（LLM）的发展，模型的通用能力逐渐增强。通常，这些模型通过混合各种任务数据进行训练，并采用多阶段训练策略。

通过阅读上面的论文，发现embedding模型也在朝这个方向发展，开始采用混合多种任务数据的训练方式。然而，由于嵌入模型的特殊性，不同任务通常使用不同的损失函数。并且根据数据的质量或其与下游任务的契合程度，将数据分为两批，进行多阶段训练。

在写这篇文章的时候，突然想到一个问题：知乎文章和视频号是不是很类似，都是选题、调研、整理、归纳的过程。而这整个过程是特别考验一个人的学习能力的，如何快速找到核心点，学明白，然后整理成别人很容易理解的样子？有时候并不是一时半会可以掌握的。