Hugging Face 使用 10 亿训练对开发的句子嵌入模型

TL;DR

Hugging Face 通过在多达 10 亿个句子对的大规模语料库上进行训练，开发出了最先进的通用句子嵌入模型。这是利用 JAX/Flax 框架和 Google TPU v3-8 基础设施，通过优化大规模对比学习实现的。

句子嵌入模型将句子映射为实数向量，以捕捉语义含义，用于聚类、文本挖掘和问答等应用。由于所有可能句子的集合是无限的，这些模型使用组合模块——通常是 Transformer 后接一个针对上下文词向量的池化操作——来计算最终的表示。

这些模型是使用一种被称为多负样本排序损失（也称为 InfoNCE 或 NTXentLoss）的对比训练方法进行训练的。这种方法使用批次内负样本（in-batch negatives）来优化嵌入空间：

数据集组成：训练集由含义接近的对 $(a_i, p_i)$ 组成（例如，查询-回答对、重复问题或被引用的论文标题）。
目标：训练模型将正样本对 $(a_i, p_i)$ 映射到接近的向量，同时将不匹配的对 $(a_i, p_j)$（其中 $i \neq j$）推向较远的向量。
相似度函数：模型计算批次中所有对之间的相似度矩阵。使用的相似度函数是余弦相似度（Cosine-Similarity）或点积（Dot-Product）。
- 余弦相似度：归一化向量确保最高相似度是与自身向量的相似度 (1)，并且与欧几里得距离成正比，使其能够与 k-means 聚类兼容。
- 点积：在使用某些近似最近邻方法时可能更快，但无法与 k-means 聚类配合使用，并且允许其他向量比向量自身与自身的相似度具有更高的点积。

为了防止分值差异变得过小，会对相似度分数应用一个缩放因子 $C$（通常 $C=20$）：$sim_{scaled}(a, b) = C * sim(a, b)$。

批次组成对于对比学习的性能至关重要。Hugging Face 确定了三个提高质量的主要杠杆：

批次大小 (Batch Size)：较大的批次大小通常与更好的模型性能相关。
难负样本 (Hard Negatives)：包含在语义上接近正样本 $p_i$ 但并非正确匹配的样本 $p_j$（例如，“What is the capital of France?” 与 “What is the capital of the US?”），这会迫使模型学习更精确的语义区分。
跨数据集批次 (Cross-Dataset Batches)：通过在单个批次中混合至少两个不同的数据集，模型可以学习不同主题之间的全局结构，而不仅仅是单个主题内的局部结构。

该项目利用了 7 个 TPU v3-8 来处理大规模对比学习所需的矩阵乘法。训练数据由拼接后的数据集组成，总计多达 10 亿个句子对。

Hugging Face 训练了 20 个基于 Mini-LM、RoBERTa、DistilBERT 和 MPNet 等架构的通用 Sentence Transformer 模型，在多个通用句子相似度评估任务上取得了最先进 (SOTA) 的结果。

除了模型之外，还发布了八个专门用于句子相似度、问答和性别评估的数据集。这些嵌入模型可以实现多种实际应用：