使用 Sentence Transformers 訓練與微調嵌入模型
Hugging Face 詳細說明了一個使用 Sentence Transformers 套件訓練與微調嵌入模型的精簡工作流程。此框架允許開發者將通用嵌入模型調整為符合特定相似度概念的需求,適用於檢索增強生成(RAG)、語意搜尋與同義句挖掘等任務。
微調對於任務特定相似度的必要性
微調至關重要,因為不同的應用需要不同的「相似度」定義。例如,兩則關於不同公司的新聞標題(例如 Apple 與 NVIDIA)在新聞分類模型中可能被視為相似,因為它們都屬於「科技」類別;但在語意文本相似度或檢索模型中必須被視為不相似,因為它們描述的是不同的事件。
核心訓練組件
訓練 Sentence Transformer 模型涉及五個主要組件:
- Dataset:訓練與評估資料,通常以
datasets.Dataset或datasets.DatasetDict實例載入。 - Loss Function:根據可用資料與目標任務量化模型效能並指導最佳化的函式。
- Training Arguments:透過
SentenceTransformersTrainingArguments提供的可選參數,用於控制訓練效率、追蹤與除錯。 - Evaluator:可選工具,用於在訓練前、訓練中或訓練後使用具體指標評估模型效能。
- 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 可以提供評估損失,但專用的 evaluator 能提供針對任務的指標。可用的 evaluator 包括:
- EmbeddingSimilarityEvaluator:針對具有相似度分數的配對(例如使用 STSb 基準)。
- TripletEvaluator:針對
(anchor, positive, negative)配對(例如使用 AllNLI 資料集)。 - InformationRetrievalEvaluator:針對查詢、語料庫與相關文件。
- BinaryClassificationEvaluator:針對具有類別標籤的配對。
可將多個 evaluator 結合成單一的 SequentialEvaluator,在訓練期間同時追蹤多種指標。
進階訓練工作流程
多資料集訓練
高效能模型通常需要同時在多個資料集上訓練。SentenceTransformerTrainer 支援此需求,接受一個資料集字典以及相對應的損失函式字典,允許在同一次訓練中對不同資料集套用不同的損失函式。
可透過 MultiDatasetBatchSamplers 從多個資料集抽樣,使用以下兩種策略:
- ROUND_ROBIN:從每個資料集等量抽樣,直到其中一個資料集耗盡。
- PROPORTIONAL:依照每個資料集的大小比例抽樣,確保所有樣本皆被使用。
轉向 SentenceTransformerTrainer
隨著 Sentence Transformers v3.0 的發布,傳統的 SentenceTransformer.fit 方法現在在內部使用 SentenceTransformerTrainer。雖然舊版程式碼仍可運作,Hugging Face 建議採用新的 trainer 方式,以利用多 GPU 訓練與改進的損失記錄等進階功能。
效能範例
在提供的範例中,使用 MultipleNegativesRankingLoss 在 AllNLI 三元組上微調 microsoft/mpnet-base 模型,取得顯著的效能提升。基礎模型在開發集上的得分為 68.32%,而微調後的模型在開發集上達到 90.04%,在測試集上達到 91.5%,以餘弦相似度的三元組準確率衡量。