使用 Sentence Transformers 进行嵌入模型的训练与微调

Hugging Face 详细介绍了一套使用 Sentence Transformers 库进行嵌入模型训练与微调的简化工作流。该框架使开发者能够将通用嵌入模型适配到任务所需的特定相似度概念，例如检索增强生成（RAG）、语义搜索和同义句挖掘。

微调对任务特定相似度的必要性

微调至关重要，因为不同的应用对“相似度”有不同的定义。例如，两条关于不同公司的新闻标题（如 Apple 和 NVIDIA）在新闻分类模型中可能被视为相似，因为它们都属于“科技”类别，但在语义文本相似度或检索模型中必须被视为不相似，因为它们描述的是不同的事件。

核心训练组件

训练 Sentence Transformer 模型涉及五个主要组件：

  1. Dataset：训练和评估数据，通常以 datasets.Dataset 或 datasets.DatasetDict 实例加载。
  2. Loss Function：一种衡量模型性能并根据可用数据和目标任务指导优化的函数。
  3. Training Arguments：通过 SentenceTransformersTrainingArguments 提供的可选参数，用于控制训练效率、跟踪和调试。
  4. Evaluator：可选工具，用于在训练前、训练中或训练后使用具体指标评估模型性能。
  5. Trainer：SentenceTransformerTrainer，将模型、数据集、损失函数及其他组件整合在一起。

数据集要求与格式化

数据集可以从 Hugging Face Hub 获取（通常带有 sentence-transformers 标签），也可以本地加载 CSV、JSON、Parquet、Arrow 或 SQL 格式的文件。为确保与所选损失函数兼容，数据集必须遵循特定的格式规则：

  • Labels：如果损失函数需要标签，数据集必须包含名为 label 或 score 的列。
  • Inputs：除标签列外的所有列均视为输入。这些列的数量和顺序必须符合损失函数的要求（例如 (anchor, positive, negative) 三元组格式）。

损失函数与训练参数

损失函数在模型训练时进行初始化。损失函数的选择取决于可用的数据（例如，针对带有浮点相似度分数的对使用 CoSENTLoss）。

可以使用 SentenceTransformersTrainingArguments 调整训练性能，其中包括 num_train_epochs、per_device_train_batch_size、warmup_ratio 以及硬件相关设置如 fp16 或 bf16。对于使用“批内负例”的损失函数，建议使用 batch_sampler=BatchSamplers.NO_DUPLICATES 参数。

模型评估

虽然 trainer 可以提供评估损失，但专用评估器能够提供任务特定的指标。可用的评估器包括：

  • EmbeddingSimilarityEvaluator：用于带有相似度分数的对（例如使用 STSb 基准）。
  • TripletEvaluator：用于 (anchor, positive, negative) 三元组（例如使用 AllNLI 数据集）。
  • InformationRetrievalEvaluator：用于查询、语料库和相关文档的评估。
  • BinaryClassificationEvaluator：用于带有类别标签的对。

多个评估器可以组合成一个 SequentialEvaluator，在训练期间同时跟踪多种指标。

高级训练工作流

多数据集训练

高性能模型通常需要同时在多个数据集上进行训练。SentenceTransformerTrainer 通过接受数据集字典和相应的损失函数字典来支持此功能。这使得在同一次训练中可以对不同的数据集使用不同的损失函数。

可以通过 MultiDatasetBatchSamplers 使用两种策略从多个数据集进行抽样：

  • ROUND_ROBIN：从每个数据集等量抽样，直至某个数据集耗尽。
  • PROPORTIONAL：按每个数据集的大小比例抽样，确保所有样本都被使用。

向 SentenceTransformerTrainer 的迁移

随着 Sentence Transformers v3.0 的发布，传统的 SentenceTransformer.fit 方法现在在内部使用 SentenceTransformerTrainer。虽然旧代码仍可运行，Hugging Face 建议采用新的 trainer 方法，以利用多 GPU 训练和改进的损失日志等高级特性。

性能示例

在提供的示例中，对 microsoft/mpnet-base 模型在 AllNLI 三元组上使用 MultipleNegativesRankingLoss 进行微调，显著提升了性能。基础模型在开发集上的得分为 68.32%，而微调后的模型在开发集上达到 90.04%，在测试集上达到 91.5%，这些指标均基于使用余弦相似度的三元组准确率。

Sources