BGE M3-Embedding來自BAAI和中國科學技術大學，是BAAI開源的模型。相關論文在https://arxiv.org/abs/2402.03216，論文提出了一種新的embedding模型，稱爲M3-Embedding，它在多語言性（Multi-Linguality）、多功能性（Multi-Functionality）和多粒度性（Multi-Granularity）方面表現出色。M3-Embedding支持超過100種工作語言，支持8192長度的輸入文本，同時支持密集檢索（Dense Retrieval）、多向量檢索（Multi-Vector Retrieval）和稀疏檢索（Sparse Retrieval），爲現實世界中的信息檢索（IR）應用提供了統一的模型基礎，通過這幾種檢索方式的組合，取得了良好的混合召回效果。
我們可以查看官方與openai模型的對比，整體來看，採用三種方式聯合檢索的BGE-M3(ALL)在三項評測中全面領先，而 BGE-M3(Dense)稠密檢索在多語言、跨語言檢索中具有明顯優勢。
BGE-M3 模型亮點
1. 多語言（Multi-Linguality），訓練集包含100+種以上語言
2. 多功能（Multi-Functionality），支持稠密檢索（Dense Retrieval），還支持稀疏檢索（Sparse Retrieval）與多向量檢索（Multi-vector Retrieval）
3. 多粒度（Multi-Granularity） BGE-M3目前可以處理最大長度爲8192 的輸入文本，支持“句子”、“段落”、“篇章”、“文檔”等不同粒度的輸入文本
BGE-M3 訓練數據
M3-Embedding模型的訓練數據組成是模型效果較好的一個關鍵創新點，因爲它旨在支持多語言性、多功能性和多粒度性，訓練數據分爲三部分：
- 無監督數據（Unsupervised Data）：
- 從大量多語言語料庫中提取未經標記的文本數據，這些語料庫包括Wikipedia、S2ORC、xP3、mC4和CC-News等。
- 這些數據通過提取豐富的語義結構（例如標題-正文、標題-摘要、指令-輸出等）來豐富模型的語義理解能力。
- 無監督數據的規模達到了1.2億文本對，覆蓋了194種語言和2655種跨語言對應關係。
- 微調數據（Fine-tuning Data）：
- 來自標記語料庫的高質量數據，包括英語、中文和其他語言的數據集。
- 例如，英語數據集包括HotpotQA、TriviaQA、NQ、MS MARCO等，而中文數據集包括DuReader、T2-Ranking、NLI-zh等。
- 這些數據集用於進一步微調模型，以提高其在特定任務和語言上的性能。
- 合成數據（Synthetic Data）：
- 爲了解決長文檔檢索任務中數據不足的問題，研究者們生成了額外的多語言微調數據（稱爲MultiLongDoc）。
- 通過從Wikipedia和MC4數據集中隨機選擇長文章，並從中隨機選擇段落，然後使用GPT-3.5生成基於這些段落的問題，生成的問題和所選文章構成新的文本對，增加了訓練數據的多樣性和覆蓋範圍。
這些訓練數據的創新之處在於：
- 多語言覆蓋：M3-Embedding通過大規模的多語言無監督數據，學習不同語言之間的共同語義空間，從而支持多語言檢索和跨語言檢索。
- 數據多樣性：通過結合無監督數據、微調數據和合成數據，M3-Embedding能夠捕捉到不同類型和長度的文本數據的語義信息，從而提高模型對不同輸入粒度的處理能力。
- 高質量數據整合：通過精心篩選和整合不同來源的數據，M3-Embedding確保了訓練數據的高質量，這對於模型學習有效的文本嵌入至關重要。
通過這種創新的訓練數據組成，M3-Embedding能夠有效地學習並支持超過100種語言的文本嵌入，同時處理從短句到長達8192個詞符的長文檔，實現了在多語言、多功能和多粒度方面的突破。
從業界來看，M3-Embedding、E5-mistral-7b，都是利用GPT這樣的LLM來合成了大量多語言數據，這個應該是後續的主流方案。
BGE-M3 混合檢索
M3-Embedding統一了嵌入模型的三種常見檢索功能，即密集檢索（Dense retrieval）、詞彙（稀疏）檢索（Lexical retrieval）和多向量檢索（Multi-vector retrieval）。以下是這些方法的公式化描述：
- 密集檢索（Dense retrieval）：輸入查詢q被轉換爲基於文本編碼器的隱藏狀態Hq，使用特殊標記“[CLS]”的歸一化隱藏狀態來表示查詢：$e_q = \text{norm}(H_q[0])$。類似地，我們可以獲取段落p的嵌入表示爲 $e_p = \text{norm}(H_p[0])$。查詢和段落之間的相關性得分通過兩個嵌入向量 $e_p$ 和 $e_q$的內積來度量：$s_{\text{dense}} \leftarrow \langle e_p, e_q \rangle$。
def dense_embedding(self, hidden_state, mask):
if self.sentence_pooling_method == 'cls':
return hidden_state[:, 0]
elif self.sentence_pooling_method == 'mean':
s = torch.sum(hidden_state * mask.unsqueeze(-1).float(), dim=1)
d = mask.sum(axis=1, keepdim=True).float()
return s / d
- 詞彙檢索（Lexical Retrieval）：輸出嵌入還被用來估計每個詞項的重要性，以促進詞彙檢索。對於查詢中的每個詞項t（在我們的工作中，詞項對應於一個標記），詞項權重被計算爲 $w_{qt} \leftarrow \text{Relu}(W_{\text{lex}} H_q[i])$，其中 $W_{\text{lex}} \in \mathbb{R}^{d \times 1}$ 是將隱藏狀態映射到一個實數的矩陣。如果詞項t在查詢中出現多次，我們只保留其最大權重。我們以相同的方式計算段落中每個詞項的權重。基於估計的詞項權重，查詢和段落之間的相關性得分通過查詢和段落中共同出現的詞項（表示爲$q ∩ p$）的聯合重要性來計算：$s_{\text{lex}} \leftarrow \sum_{t \in q \cap p}(w_{qt} \cdot w_{pt})$。
def sparse_embedding(self, hidden_state, input_ids, return_embedding: bool = True):
# sparse_linear 線性層= torch.nn.Linear(in_features=self.model.config.hidden_size, out_features=1)
# 通過relu計算token weight
token_weights = torch.relu(self.sparse_linear(hidden_state))
if not return_embedding: return token_weights
# 形狀爲(input_ids.size(0), input_ids.size(1), self.vocab_size)的零張量
sparse_embedding = torch.zeros(input_ids.size(0), input_ids.size(1), self.vocab_size,
dtype=token_weights.dtype,
device=token_weights.device)
# 將token_weights中的值分散scatter到sparse_embedding的相應位置，索引位置根據input_ids提供
sparse_embedding = torch.scatter(sparse_embedding, dim=-1, index=input_ids.unsqueeze(-1), src=token_weights)
# CLS，PAD 等無用token
unused_tokens = [self.tokenizer.cls_token_id, self.tokenizer.eos_token_id, self.tokenizer.pad_token_id,
self.tokenizer.unk_token_id]
sparse_embedding = torch.max(sparse_embedding, dim=1).values
# 無用token weight設置爲0
sparse_embedding[:, unused_tokens] *= 0.
return sparse_embedding
- 多向量檢索（Multi-Vector Retrieval）：作爲密集檢索的擴展，多向量方法利用整個輸出嵌入來表示查詢和段落：$E_q = \text{norm}(W_{\text{mul}} H_q), E_p = \text{norm}(W_{\text{mul}} H_p)$，其中 $W_{\text{mul}} \in \mathbb{R}^{d \times d}$ 是可學習的投影矩陣。按照ColBERT（Khattab和Zaharia, 2020）的方法，使用後期交互來計算細粒度的相關性得分：$s_{\text{mul}} \leftarrow \frac{1}{N} \sum_{i=1}^{N} \max_{j=1}^{M} E_q[i] \cdot E_p[j]$；N和M分別是查詢和段落的長度。
由於嵌入模型的多功能性，檢索過程可以在混合過程中進行。首先，可以通過每種方法單獨檢索候選結果（由於其高成本，可以免去多向量方法的這一步）。最終的檢索結果是根據集成的相關性得分重新排序的：$s_{\text{rank}} \leftarrow s_{\text{dense}} + s_{\text{lex}} + s_{\text{mul}}$。
BGE-M3訓練方式和創新點
BGE-M3模型訓練分爲三個階段：
-
1）RetroMAE預訓練，在105種語言的網頁數據和wiki數據上進行，提供一個可以支持8192長度和麪向表示任務的基座模型；
-
2）無監督對比學習，在194種單語言和1390種翻譯對數據共1.1B的文本對上進行的大規模對比學習；
-
3）多檢索方式統一優化，在高質量多樣化的數據上進行多功能檢索優化，使模型具備多種檢索能力。
-
其中，一些重要的關鍵技術如下：
1. 自學習蒸餾
人類可以利用多種不同的方式計算結果，矯正誤差。模型也可以，通過聯合多種檢索方式的輸出，可以取得比單檢索模式更好的效果。因此，BGE-M3使用了一種自激勵蒸餾方法來提高檢索性能。具體來說，合併三種檢索模式的輸出，得到新的文本相似度分數，將其作爲激勵信號，讓各單模式學習該信號，以提高單檢索模式的效果。
2. 訓練效率優化
通過根據長度對文本數據進行分組，確保一個batch內文本長度相對相似，從而減少填充。爲了減少文本建模時的顯存消耗，將一批數據分成多個小批。對於每個小批，利用模型編碼文本，收集輸出的向量同時丟棄所有前向傳播中的中間狀態，最後彙總向量計算損失，可以顯著增加訓練的batch size。
3. 長文本優化
BGE-M3提出了一種簡單而有效的方法：MCLS(Multiple CLS)來增強模型的能力，而無需對長文本進行微調。
MCLS方法旨在利用多個CLS令牌來聯合捕獲長文本的語義。爲每個固定數量的令牌插入一個cls令牌，每個cls令牌可以從相鄰的令牌獲取語義信息，最後通過對所有cls令牌的最後隱藏狀態求平均值來獲得最終的文本嵌入。
BGE-M3 實驗結果
多語言檢索任務，稀疏檢索(Sparse)大幅超過了傳統的稀疏匹配算法BM25。多向量檢索(multi-vector)則獲得了三種檢索方式中的最佳效果。
跨語言檢索能力（MKQA）
BGE-M3在跨語言任務上依然具備最佳的檢索效果。稀疏檢索並不擅長應對跨語言檢索這種詞彙重合度很小的場景。因此，稀疏檢索的自身效果以及與其他方法混搭所帶來的收益相對較小。
長文檔檢索能力 (MLRB: Multi-Lingual Long Retrieval Benchmark)
BGE-M3可以支持長達8192的輸入文檔，從實驗結果可以觀察到，稀疏檢索（Sparse）的效果要顯著高於稠密檢索（Dense），這說明關鍵詞信息對於長文檔檢索極爲重要。
BGE-M3 模型微調
需要先安裝，
- with pip
pip install -U FlagEmbedding
- from source
git clone https://github.com/FlagOpen/FlagEmbedding.git
cd FlagEmbedding
pip install -e .
模型微調的數據集格式是json line格式文件，json格式如下：
{"query": str, "pos": List[str], "neg":List[str]}
query 是查詢，pos 是正文本列表，neg 是負文本列表。
模型訓練：
torchrun --nproc_per_node {number of gpus} \
-m FlagEmbedding.BGE_M3.run \
--output_dir {path to save model} \
--model_name_or_path BAAI/bge-m3 \
--train_data ./toy_train_data \
--learning_rate 1e-5 \
--fp16 \
--num_train_epochs 5 \
--per_device_train_batch_size {large batch size; set 1 for toy data} \
--dataloader_drop_last True \
--normlized True \
--temperature 0.02 \
--query_max_len 64 \
--passage_max_len 256 \
--train_group_size 2 \
--negatives_cross_device \
--logging_steps 10 \
--same_task_within_batch True \
--unified_finetuning True \
--use_self_distill True
參考文獻
- 1 BGE M3-Embedding: Multi-Lingual, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation
https://github.com/FlagOpen/FlagEmbedding/blob/master/FlagEmbedding/BGE_M3/BGE_M3.pdf