使用 Sentence Transformers 訓練與微調多向量嵌入模型

Sentence Transformers 是一個 Python 函式庫，用於使用和訓練嵌入模型及重排序模型，應用範圍廣泛，例如檢索增強生成、語義搜尋和語義文本相似度等。其 v6.0 更新引入了第四種模型類型：MultiVectorEncoder，專為 ColBERT 風格的延遲互動檢索設計，並提供完整的訓練方法。

在這篇部落格文章中，我將展示如何使用它來微調一個多向量模型，使其在您的資料上表現優於通用檢索器。這種方法也可以從頭開始訓練強大的新型多向量模型。以下所有內容都可以在 `pip install -U "sentence-transformers[train]"` 後運行。

微調多向量模型涉及多個組成部分：模型本身、資料集、損失函數、訓練參數、評估器和訓練器類別。我將逐一檢視這些組成部分，並附上實用的範例，說明如何將它們用於微調強大的多向量模型。

最後，在「評估」部分，我將展示我的微調多向量編碼器/mLateOn-medical 模型，這個模型與這篇部落格文章同時在單一 RTX 3090 上訓練了 14.5 小時，它在我的醫療檢索評估中，輕鬆超越了我能找到的所有通用檢索模型，無論是密集、稀疏、詞彙或多向量模型。

如果您對微調密集嵌入模型、稀疏嵌入模型或重排序器感興趣，那麼請考慮閱讀我之前的部落格文章：《訓練與微調嵌入模型》、《訓練與微調稀疏嵌入模型》以及《訓練與微調重排序器模型》。

這篇部落格文章是關於訓練多向量模型的。如果您想學習如何使用它們，從載入、編碼到在向量資料庫中建立索引，請參閱配套的部落格文章：《使用 Sentence Transformers 的多向量（延遲互動）嵌入模型》。

什麼是多向量模型？

密集嵌入模型將整個文本壓縮成單一向量，而相似度則是兩個摘要之間的點積。多向量模型（也稱為延遲互動或 ColBERT 風格模型）則跳過這種壓縮。它為每個 token 保留一個小向量，並使用 MaxSim 運算子來評分查詢與文件之間的相似度，其中每個查詢 token 找到其最佳匹配的文件 token，然後將分數加總。

Token 級別的匹配精確保留了單一向量必須平均掉的細粒度訊號，這通常意味著更強大的檢索能力，但代價是更大的索引。配套的《多向量嵌入模型》部落格文章詳細介紹了其架構、編碼、評分和索引，因此本節我將簡短說明，直接進入訓練部分。

為何要微調？

微調多向量模型能顯著提升其在特定領域的檢索效能：網路搜尋、法律文件探索、程式碼搜尋和科學文獻回顧等領域的詞彙、查詢風格和相關性概念都各不相同。由於查詢和文件是逐 token 匹配的，多向量模型能捕捉到單向量模型傾向於平均掉的細粒度領域訊號，即使只有少量領域內微調資料，它們也能表現得非常好。

除此之外，大多數已發布的檢索模型都是針對短篇幅內容配置的。經典的 ColBERT 檢查點會將文件截斷至 180 或 300 個 token，許多流行的密集模型則截斷至 256 或 512 個 token，因為它們的 MS MARCO 風格訓練資料很少超出這個範圍。

如果您的文件很長，這些模型在評分前會默默地丟棄大部分文件內容。在我的醫療評估中，平均篇幅為 941 個 token，我測量到這種截斷會導致高達 0.24 的 NDCG@10 損失，這遠超過模型架構之間的任何差異。當您訓練自己的模型時，您可以根據資料需求配置文件長度。

LightOn 在程式碼檢索中也遇到了同樣的動態，通用 LateOn 不夠用，他們因此訓練了 LateOn-Code。您的領域，無論是醫療、法律、金融或您公司的內部文件，都不會有官方模型。這篇部落格文章將展示如何在數小時內，使用單一消費級 GPU，自行建立這樣的模型。

訓練組成部分

訓練 MultiVectorEncoder 模型涉及以下組成部分：

模型：要微調的模型或要從頭建立的架構。

資料集：用於訓練和評估的資料。

損失函數：衡量模型效能並指導最佳化過程的函數。

訓練參數（可選）：影響訓練效能、追蹤和除錯的參數。

評估器（可選）：用於在訓練前、訓練期間或訓練後評估模型的類別。

訓練器：將所有訓練組成部分整合在一起。

讓我們仔細看看每個組成部分。

模型

多向量訓練讓您在起點上有真正的選擇，而且這比您想像的更重要。

微調現有的多向量模型

如果您想進一步微調現有的多向量模型，您完全不必擔心架構問題：

```python

from sentence_transformers import MultiVectorEncoder

# Loading in fp32 is preferred for training if your memory can handle it

model = MultiVectorEncoder(

"lightonai/mLateOn-unsupervised",

model_kwargs={"torch_dtype": "float32"},

processor_kwargs={"model_max_length": 8192}, # the tokenizer-level token limit

)

```

檢查點會自帶其配方：其查詢和文件標記 token、其投影頭、其評分跳過列表。對於微調，您通常希望保留所有這些，只更改資料所需的部分。首先要檢查的是長度配置，因為許多已發布的檢查點將文件限制在 180 到 512 個 token（請參閱「為何要微調？」）

，而我的醫療段落長達 1,400 個 token。mLateOn 系列已經支援骨幹模型的完整 8192 token 上下文，但如果您的起始檢查點帶有上限，請將其解除：

```python

# Let the model read full documents instead of the caps it was trained with,

# e.g. GTE-ModernColBERT-v1 ships with query_length=48 and document_length=300

model[0].query_length = None

model[0].document_length = None

```

在取消每個任務的上限後，截斷會退回到分詞器的 `model_max_length`，這就是為什麼我在載入時配置該限制的原因。我還做了一個更改，添加了一個標點符號跳過列表，將標點符號 token 從文件側評分和儲存中排除。在四種消融實驗（無、標點符號、停用詞、兩者）中，它在品質上略勝一籌，並且在不影響效能的情況下，將此資料的文件索引縮小了 9.6%：

```python

import string

# model[2] is the MultiVectorMask module

model[2].skiplist_words = list(string.punctuation)

model[2].resolve_with_tokenizer(model.tokenizer) # token ids are cached, so re-resolve after changing

```

從基礎 Transformer 建立模型

您也可以將 MultiVectorEncoder 指向任何基礎 transformer，它會為您附加一個全新、隨機初始化的 token 級別投影：

```python

from sentence_transformers import MultiVectorEncoder

model = MultiVectorEncoder("answerdotai/ModernBERT-base", model_kwargs={"torch_dtype": "float32"})

# MultiVectorEncoder(

# (0): Transformer({..., 'architecture': 'ModernBertModel'})

# (1): Dense({'in_features': 768, 'out_features': 128, 'bias': False, ...})

# (2): MultiVectorMask({'skiplist_words': [], 'skiplist_tasks': ['document'], ...})

# (3): Normalize({...})

# )

```

這就是經典的 ColBERT 管線：一個 Transformer 產生上下文相關的 token 嵌入，一個 token 級別的 Dense 層將每個嵌入投影到 128 維，一個 MultiVectorMask 決定哪些 token 在評分時計數，以及一個 token 級別的 Normalize 層。

投影是隨機開始的，因此在模型有用之前需要進行訓練。有趣的是，這也適用於強大的密集嵌入骨幹模型。在我的實驗中，在 Alibaba-NLP/gte-modernbert-base 上使用一個新的投影，僅憑投影和 25k 訓練對，就能達到與現有檢查點起始點相差 0.03 的水準。

經典的 ColBERT 分詞技巧（[MASK] 查詢擴展、[Q] / [D] 前綴 token、文件長度上限、標點符號跳過列表）預設都是關閉的，並且可以配置。有關完整設定，請參閱「建立自訂模型」。值得一提的是，我在我的領域微調中測試了四種 [MASK] 查詢擴展配置，但它們都沒有產生可測量的差異，所以不必強求經典配方。

您應該選擇哪個起始點？

我在準備這篇部落格文章時直接測量了這一點，選取了六個起始點，並以相同的配方在來自 MIRIAD 的 25k 醫療問題-段落對上訓練每個模型，然後在 1,000 個保留問題上對 50,000 個段落語料庫進行評估：

起始點 零樣本 NDCG@10 25k 對後 差異

lightonai/mLateOn-unsupervised 0.9087 0.9398 +0.0311

lightonai/mLateOn 0.9277 0.9319 +0.0042

lightonai/LateOn-unsupervised 0.9026 0.9206 +0.0180

lightonai/LateOn 0.9185 0.9105 -0.0080

lightonai/GTE-ModernColBERT-v1 0.9198 0.9007 -0.0191

Fresh head on gte-modernbert-base - 0.9177 -

結果讓我感到驚訝，並且在兩個模型家族中都得到了驗證。未經監督訓練的檢查點（-unsupervised）比其已完成的同類模型更能適應新領域，儘管起點較低，但最終超越了它們。這些檢查點經過大規模對比預訓練，但尚未進行通用檢索的監督微調，因此它們帶有所有延遲互動結構，卻沒有通用微調的痕跡，而領域訓練則無需撤銷這些痕跡。相比之下，已完成的檢查點幾乎沒有變化，甚至在我嘗試的每個學習率下都出現了退步。

因此，如果您喜歡的模型家族發布了預監督檢查點，請從那裡開始。如果沒有，那麼在強大的檢索預訓練骨幹模型上使用一個新的投影是一個不錯的選擇。儘管從一個完全完成的檢查點繼續訓練感覺最自然，但對於領域適應而言，這是最弱的選項。

資料集

MultiVectorEncoderTrainer 使用 `datasets.Dataset` 或 `datasets.DatasetDict` 實例進行訓練和評估。您可以從 Hugging Face Datasets Hub 載入資料，或使用您喜歡的任何格式（例如 CSV、JSON、Parquet、Arrow 或 SQL）的本地資料。

注意：許多與 Sentence Transformers 開箱即用的公開資料集已在 Hugging Face Hub 上被標記為 `sentence-transformers`，因此您可以在 `https://huggingface.co/datasets?other=sentence-transformers` 輕鬆找到它們。考慮瀏覽這些資料集，以找到可能對您的任務、領域或語言有用的現成資料集。

Hugging Face Hub 上的資料

您可以使用 `load_dataset` 函數從 Hub 上的資料集載入資料：

```python

from datasets import load_dataset

train_dataset = load_dataset("tomaarsen/miriad-4.4M-split", split="train")

print(train_dataset)

"""

Dataset({

features: ['question', 'passage_text'],

num_rows: 4467542

})

"""

```

這就是我在這篇部落格文章中將要訓練的資料集：來自 MIRIAD 的 440 萬個醫療問題，每個問題都與包含其答案的來源段落配對（平均 941 個 token）。像這樣簡單的（查詢、相關段落）對是您自己領域最容易收集的檢索訓練資料，而且如您所見，它們就是您所需要的一切。

本地資料

您也可以使用 `load_dataset` 載入常見檔案格式的本地資料：

```python

from datasets import load_dataset

dataset = load_dataset("csv", data_files="my_"

```