使用 Sentence Transformers 訓練與微調重新排序模型

Hugging Face 詳細說明了一個使用 Sentence Transformers 套件訓練與微調重新排序模型（Cross Encoder）的框架。透過在領域特定資料上進行訓練，開發者可以打造小型且高效的重新排序模型，在專門的檢索任務中超越遠大於其規模的通用模型。

重新排序模型 vs. 嵌入模型

Cross Encoder 重新排序模型透過將文字對（例如查詢與文件）一起輸入共享的神經網路，產生單一的相關性分數，以評估其相關程度。這與 bi-encoder（嵌入模型）不同，後者會將文字獨立嵌入為向量，並透過距離度量計算相似度。

雖然 Cross Encoder 因需處理每一個可能的文字對而計算成本較高，但通常能提供更高的準確度。因此，業界對高效能搜尋的標準做法是兩階段的「檢索與重新排序」流程：先使用快速的嵌入模型進行初步檢索，再以 Cross Encoder 精煉前 k 名的結果。

核心訓練組件

訓練重新排序模型涉及五個主要組件，透過 CrossEncoderTrainer 整合：

1. 資料集

CrossEncoderTrainer 支援 datasets.Dataset 或 datasets.DatasetDict 物件。資料可以來自 Hugging Face Hub 或本機檔案（CSV、JSON、Parquet、Arrow、SQL）。為確保與所選損失函式相容，資料集必須包含名為「label」「labels」「score」或「scores」的欄位作為標籤，且其餘欄位需符合損失函式所需的輸入數量。

2. 硬負樣本挖掘

模型效能在很大程度上取決於負樣本的品質。雖然「軟負樣本」完全不相關，「硬負樣本」則是看似相關卻實際不相關的段落。Sentence Transformers 提供 mine_hard_negatives 函式來辨識這些具挑戰性的樣本，迫使模型在區分上更加精確。

3. 損失函式

損失函式根據可用資料引導最佳化過程。雖然有 LambdaLoss 或 ListNetLoss 等學習排序的損失函式，BinaryCrossEntropyLoss 仍是對標記對組而言既高效又簡單的選擇。

4. 訓練參數

可透過 CrossEncoderTrainingArguments 進行自訂，讓開發者調整學習率、批次大小、暖身比例以及精度設定（FP16/BF16）。對於使用批內負樣本的損失函式，特別建議使用 batch_sampler=BatchSamplers.NO_DUPLICATES 設定。

5. 評估器

為了在單純損失之外追蹤效能，Sentence Transformers 提供多種內建評估器：

  • CrossEncoderClassificationEvaluator：用於二元或多類別標籤。
  • CrossEncoderCorrelationEvaluator：用於相似度分數（例如使用 STSb 資料集）。
  • CrossEncoderRerankingEvaluator：使用查詢、正例與負例評估重新排序效能。
  • CrossEncoderNanoBEIREvaluator：針對英文重新排序的輕量級評估器。

多資料集訓練

CrossEncoderTrainer 支援同時在多個資料集上訓練，即使它們格式不同或需要不同的損失函式。這透過資料集字典以及可選的損失函式字典來管理。抽樣策略包括：

  • ROUND_ROBIN：從每個資料集等量抽樣，直到其中一個耗盡。
  • PROPORTIONAL：依各資料集大小比例抽樣，確保所有樣本皆被使用。

效能評估與結果

在實務應用中，我們以 ModernBERT-base 為基礎的重新排序模型，在 GooAQ 資料集的 99k 查詢-答案對上進行微調。使用 BinaryCrossEntropyLoss 及硬負樣本挖掘，最終模型（tomaarsen/reranker-ModernBERT-base-gooaq-bce）在 13 種常見開源重新排序模型中表現最佳，甚至超越規模高達其四倍的模型。

GooAQ NDCG@10 結果（前 30 名重新排序）

模型 參數 真實 NDCG@10 評估 NDCG@10
Retriever only (No reranking) - 59.12 59.12
BAAI/bge-reranker-large 560M 73.20 77.46
mixedbread-ai/mxbai-rerank-large-v2 1.54B 75.40 80.04
ModernBERT-base-gooaq-bce 150M 77.14 83.51
ModernBERT-large-gooaq-bce 396M 79.42 85.81

技術訓練技巧

  • 防止過度擬合：Cross Encoder 容易快速過擬合。使用帶有 load_best_model_at_end 與 metric_for_best_model 的評估器，以捕捉最佳效能的模型。
  • 平衡負樣本：僅依賴硬負樣本可能會削弱在較簡單任務上的表現。將隨機負樣本與硬負樣本混合使用可緩解此問題。
  • 效率：在領域特定資料上微調小型重新排序模型，不僅能提升搜尋準確度，亦可相較於使用龐大通用模型降低推論延遲。

Sources