Hugging Face 使用 10 億對訓練數據開發的句子嵌入模型

TL;DR

Hugging Face 透過在多達 10 億對句子的海量語料庫上進行訓練，開發出最先進的通用型句子嵌入模型。這是利用 JAX/Flax 框架和 Google TPU v3-8 基礎設施，來優化大規模對比學習（contrastive learning）而實現的。

訓練方法論

模型架構

句子嵌入模型將句子映射到實數向量，以捕捉語義意義，用於聚類、文本挖掘和問答等應用。由於所有可能句子的集合是無限的，這些模型使用組成模組——通常是 Transformer 後接一個針對情境化字詞向量（contextualized word vectors）的池化操作（pooling operation）——來計算最終的表示。

Multiple Negative Ranking Loss (MNRL)

這些模型是使用一種稱為 Multiple Negative Ranking Loss（也稱為 InfoNCE 或 NTXentLoss）的對比訓練方法進行訓練的。這種方法使用批次內負樣本（in-batch negatives）來優化嵌入空間：

數據集組成：訓練集由意義相近的對 $(a_i, p_i)$ 組成（例如：查詢-回答對、重複問題或被引用的論文標題）。
目標：訓練模型將正樣本對 $(a_i, p_i)$ 映射到接近的向量，同時將不匹配的對 $(a_i, p_j)$（其中 $i eq j$）推向遠離的向量。
相似度函數：模型計算批次中所有對之間的相似度矩陣。使用的相似度函數是 Cosine-Similarity 或 Dot-Product。
- Cosine-Similarity：歸一化向量確保最高相似度是與向量本身（1），且與歐幾里得距離成正比，使其與 k-means 聚類相容。
- Dot-Product：在使用某些近似最近鄰方法時可能更快，但無法與 k-means 聚類配合使用，且允許其他向量比向量與自身的相似度具有更高的點積。

為了防止分數差異變得太小，會對相似度分數應用一個縮放因子 $C$（通常 $C=20$）：$sim_{scaled}(a, b) = C * sim(a, b)$。

優化嵌入品質

批次大小與硬負樣本 (Hard Negatives)

批次組成對於對比學習的性能至關重要。Hugging Face 識別出三個主要提升品質的槓桿：

Batch Size：較大的批次大小通常與更好的模型性能相關。
Hard Negatives：包含在語義上與正樣本 $p_i$ 接近但並非正確匹配的樣本 $p_j$（例如：「法國的首都是哪裡？」與「美國的首都是哪裡？」），這會迫使模型學習更精確的語義區別。
Cross-Dataset Batches：透過在單一批次中混合至少兩個不同的數據集，模型可以學習不同主題之間的全局結構，而不僅僅是單一主題內的局部結構。

基礎設施與數據規模

該專案利用了 7 個 TPU v3-8 來處理大規模對比學習所需的矩陣乘法。訓練數據由串聯的數據集組成，總計多達 10 億對句子。

結果與應用

Hugging Face 訓練了 20 個通用型 Sentence Transformer 模型，其架構包括 Mini-LM、RoBERTa、DistilBERT 和 MPNet，在多個通用型句子相似度評估任務中達到了最先進的 (SOTA) 結果。

連同模型一起發布的還有八個用於句子相似度、問答和性別評估的專業化數據集。這些嵌入模型可以實現多種實際應用：

Sentence Similarity：使用餘弦相似度（cosine similarity）比較兩個文本的語義接近程度。
Asymmetric QA：判斷候選段落回答特定查詢的可能性。
Search and Clustering：使用點積距離（dot-product distance）檢索與查詢相近的答案。
Gender Bias Evaluation：透過比較模型在以職業為基礎的錨點文本中對性別代詞的相似度分數，來識別訓練集中的固有性別偏見。