模型描述

该项目旨在使用自监督对比学习目标在非常大的句子级数据集上训练句子嵌入模型。我们使用预训练的 roberta-large 模型，并在一个10亿个句对的数据集上进行微调。我们使用对比学习目标：给定一对句子中的一个句子，模型应该预测哪个随机抽样的另一个句子实际上是与之配对的。

我们在由Hugging Face组织的 Community week using JAX/Flax for NLP & CV 期间开发了这个模型。我们开发了这个模型作为项目 Train the Best Sentence Embedding Model Ever with 1B Training Pairs 的一部分。我们受益于高效的硬件基础设施，用于运行项目：7个 TPUs v3-8，以及来自Google的Flax、JAX和Cloud团队成员的关于高效深度学习框架的干预。

预期用途

我们的模型可用作句子编码器。给定一个输入句子，它输出一个捕捉句子语义信息的向量。该句向量可用于信息检索、聚类或句子相似性任务。

如何使用

以下是使用此模型获取给定文本特征的方法，使用 SentenceTransformers 库：

from sentence_transformers import SentenceTransformer

model = SentenceTransformer('flax-sentence-embeddings/all_datasets_v3_roberta-large')
text = "Replace me by any text you'd like."
text_embbedding = model.encode(text)
# array([-0.01559514,  0.04046123,  0.1317083 ,  0.00085931,  0.04585106,
#        -0.05607086,  0.0138078 ,  0.03569756,  0.01420381,  0.04266302 ...],
#        dtype=float32)

训练过程

预训练

我们使用预训练的 roberta-large 模型。有关预训练过程的详细信息，请参阅模型卡。

微调

我们使用对比目标对模型进行微调。形式上，我们从批次中的每对句子计算余弦相似度。然后，将其与真实的句子对进行交叉熵损失比较。

超参数

我们在 TPU v3-8 上训练了模型。我们使用批大小为1024（每个TPU核心128个）在540,000个步骤中训练模型。我们使用了500次学习率预热。序列长度限制为128个标记。我们使用了学习率为2e-5的AdamW优化器。完整的训练脚本可在当前的存储库中访问。

训练数据

我们使用多个数据集的连接来微调我们的模型。句对的总数超过10亿个句子。我们根据数据配置文件（data_config.json）中详细说明的加权概率对每个数据集进行采样。