从零训练文本嵌入模型:Arctic Embed 方法解析
嵌入模型(Embedding models)在无需额外调优的情况下,能够提供准确的检索性能,这使得它们在搜索和检索增强生成(RAG)工作负载中备受青睐。与传统的关键词搜索不同,嵌入模型超越了词汇重叠的限制进行信息编码,将文本映射到高维向量空间,使得语义相似的文本在空间中距离更近。
由于这些模型的实用性及其广泛采用,开放源代码和研究社区不断推出越来越强大的文本嵌入模型,如 E5、GTE 和 Jina。这些工作的快速实验和改进,部分归功于大型开放评估基准的支持,如 MSMARCO、BEIR 和 MTEB。这些基准测试平台为研究人员提供了可靠的评估环境,促进了嵌入模型的持续发展和优化。
本文基于 Arctic Embed 项目的实践,提出了一组消融实验,表明在训练过程中数据采样和负采样方法比扩大数据规模和批量大小更能显著改善检索质量,而此前的工作主要集中在后者。此外,本文还介绍了一种基于挖掘的困难负样本的创新查询生成技术,发现其比同时生成查询和负样本的直接生成方法更有效。
Arctic 嵌入模型架构
在 Arctic-embed 的开发中,我们旨在从文献中公认的最佳实践出发,从头开始训练一个嵌入模型。与 E5、BGE、GTE、Jina 和 Nomic 等先前工作一致,我们进行了两轮训练,使用了两种不同类型的数据集。
网络结构
我们基于不同规模的 BERT 类模型进行了训练。我们的 m 和 l 模型采用标准的 BERT 架构(分别为 BERT base 和 BERT large)。对于较小的模型(xs 和 s),我们选择了 MiniLMv2 架构的变体,采用了 Nomic BERT 架构。这种分层设计允许我们在计算资源受限的场景下也能获得良好的性能表现。
池化层
在架构上,我们没有对任何基础模型进行修改,甚至没有进行常见的添加池化层的操作。此外,我们使用 [CLS] 标记的最终隐藏状态作为嵌入向量,而不是 E5、GTE 和 Nomic 中使用的平均池化策略。这一选择与 BGE 架构一致,实验表明 [CLS] token 往往能更好地捕捉句子的整体语义表示。
训练数据集构建
在创建我们的训练数据集时,我们从大型语言模型(LLMs)领域汲取灵感,并利用了受到 RefinedWeb、C4、Gopher 和 TogetherAI 启发的过滤方法。
首先,对于嘈杂的原始数据源,如网络搜索,我们使用 trafilatura 解析结构化的网页文档。在解析过程中,我们计算用于质量过滤的自定义信号。具体来说,对于正数据对清理,我们需要确保:
- 每对文本都具有良好的质量(语言过滤器、文本质量过滤器)。
- 文本对(查询,文档)在意义上相似(一致性过滤器)。
对于质量过滤,我们利用了一系列与 Snowflake 的 Arctic 模型训练手册中详细介绍的类似的过滤器。对于一致性过滤,我们采用了低保真、高吞吐量的成对相似性一致性过滤器 — 使用 fastText 的 word2vec 模型进行句子相似性计算(可以在 CPU 上便宜地运行)。我们不将这些嵌入信号视为明确的质量标签,而是采用保守的阈值(最低允许相似度为 0.3),并用它们来过滤掉无关的示例。
此外,在此步骤中,我们将长序列截断为 512 个单词。正如我们观察到的那样,网络语料库中的查询通常在文档的开头得到回答。这不仅在计算上是浪费的,而且即使是 word2vec 嵌入中捕获的含义也会因为从后面的不相关单词中平均向量而被稀释。
数据集混合与采样
由于不同数据集的大小、一致性、难度和学习动态的不同,简单地将所有可用数据集连接在一起被证明是一种次优策略,特别是在微调阶段。相反,我们进行了孤立实验,以了解每个数据集对微调性能的影响。然后,我们根据这些实验中它们的相对性能选择并组合数据集。
预训练数据集
我们的大规模预训练数据集包括了 3.08 亿个查询 - 文档对(从大约 20 亿个文档中筛选),其中 71% 是与查询或标题配对的网络搜索文档。除了网络搜索数据外,文本对集还包括了来自常见抓取源的 PAQ、StackExchange 标题 - 正文和标题 - 正文网络文档对,以及 S2ORC 标题 - 摘要对。
微调数据集
我们的微调数据集由我们的网络搜索数据与几个公共数据集(HotpotQA、NQ、Fever 和 StackExchange 标题 - 正文)组合而成,然后通过以下章节详细描述的合成挖掘策略进行进一步扩展。这种混合明显省略了其他嵌入模型使用的几个流行公共数据集,因为我们观察到正样本一致性和负样本难度水平。这些发现不太有用的数据集包括 NLI、MEDI、WikiAnswers 和 SQuAD。
经验上,我们观察到,在微调阶段,质量比数量更重要,过量的低质量数据可能会导致模型质量降低。
合成数据
与预训练中使用的大量网络规模数据相比,适用于微调的高质量示例更为稀缺。然而与这些先前方法不同,我们发现向我们的 LLM 输入中添加负面文档对于确定查询生成至关重要。此外,我们选择只生成合成查询,而不是合成负样本,因为我们发现 LLM 不容易生成与从现有文档语料库中挖掘的同样高质量的相关负样本。
图 4 展示了这一方法的实施方式——由算法 2 的变体生成的两个数据集导致了接近原始 HotpotQA 所提供的得分增加。


