RAG(五)BGE-M3,最流行的開源text embedding模型
項(xiàng)目地址:https://huggingface.co/BAAI/bge-m3
首先說明下為什么需要text embedding?
Text Embedding 是一種將文本數(shù)據(jù)映射到高維向量空間的技術(shù),這些向量能夠捕捉文本的語義信息。通過將文本嵌入到向量空間中,我們可以利用向量之間的距離或相似性來衡量文本之間的語義相關(guān)性。
在信息檢索領(lǐng)域,Text Embedding 的重要性尤為突出。傳統(tǒng)的基于關(guān)鍵詞匹配的檢索方法往往無法準(zhǔn)確理解用戶的真實(shí)意圖,也無法處理語義相似但用詞不同的查詢。而通過 Text Embedding,我們可以將查詢和文檔都嵌入到同一個(gè)語義空間中,從而能夠更準(zhǔn)確地檢索出與用戶查詢語義相關(guān)的文檔,而不僅僅是包含相同關(guān)鍵詞的文檔。
此外,隨著多語言內(nèi)容的不斷增加,跨語言檢索的需求也日益增長。Text Embedding 能夠?qū)⒉煌Z言的文本映射到一個(gè)統(tǒng)一的語義空間中,使得跨語言檢索成為可能。
智源發(fā)布的BGE-M3在多語言、多功能和多粒度方面展現(xiàn)出前所未有的通用性,能夠支持超過100種語言的語義檢索,并且可以處理從句子到長文檔的各種輸入。下面來詳細(xì)看下:
1、方法介紹
M3-Embedding 實(shí)現(xiàn)了三方面的多功能性,支持多種語言,并能夠處理不同粒度的輸入數(shù)據(jù)。此外,它還統(tǒng)一了文本嵌入的常見檢索功能。形式上,給定任意語言 x 中的查詢 q,它能夠從語料庫中檢索出語言 y 中的文檔
。
數(shù)據(jù)整理
BGE-M3-Embedding 需要一個(gè)大規(guī)模且多樣化的多語言數(shù)據(jù)集。為了構(gòu)建這樣的數(shù)據(jù)集,從三個(gè)來源進(jìn)行了全面的數(shù)據(jù)收集:
- 無監(jiān)督數(shù)據(jù):來源包括 Wikipedia、S2ORC、xP3、mC4、CC-News 和 MTP 等多語言語料庫。從這些語料庫中提取豐富的語義結(jié)構(gòu),例如標(biāo)題-正文、標(biāo)題-摘要、指令-輸出等。為了學(xué)習(xí)跨語言語義匹配的統(tǒng)一嵌入空間,引入了平行句子數(shù)據(jù),來自 NLLB 和 CCMatrix 兩個(gè)翻譯數(shù)據(jù)集。總共整理出1.2 億個(gè)文本對(duì),涵蓋 194 種語言和 2655 種跨語言對(duì)應(yīng)關(guān)系。
- 微調(diào)數(shù)據(jù):
英語:整合了 8 個(gè)數(shù)據(jù)集,包括 HotpotQA、TriviaQA、NQ、MS MARCO、COLIEE、PubMedQA、SQuAD 和 SimCSE 的 NLI 數(shù)據(jù)。
中文:整合了 7 個(gè)數(shù)據(jù)集,包括 DuReader、mMARCO-ZH、、LawGPT、CMedQAv2、
和 LeCaRDv2。
其他語言:利用 Mr. TyDi 和 MIRACL 的訓(xùn)練數(shù)據(jù)。 - 合成數(shù)據(jù):為緩解長文檔檢索任務(wù)的短缺,生成了額外的多語言微調(diào)數(shù)據(jù)(稱為 MultiLongDoc)。從 Wikipedia、Wudao 和 mC4 數(shù)據(jù)集中抽取長篇文章,并從中隨機(jī)選擇段落。使用 GPT-3.5 根據(jù)這些段落生成問題,生成的問題和抽取的文章構(gòu)成新的文本對(duì)。
混合檢索
BGE-M3-Embedding 統(tǒng)一密集檢索、詞匯(稀疏)檢索和多向量檢索的具體實(shí)現(xiàn)如下:
- 密集檢索:輸入查詢q通過文本編碼器轉(zhuǎn)換為隱藏狀態(tài)
,使用特殊標(biāo)記 "[CLS]" 的歸一化隱藏狀態(tài)表示查詢:
。同樣,段落p的嵌入為
。查詢和段落之間的相關(guān)性得分通過內(nèi)積計(jì)算:
。
- 詞匯檢索:輸出嵌入用于估計(jì)每個(gè)詞項(xiàng)的重要性,以促進(jìn)詞匯檢索。對(duì)于查詢中的每個(gè)詞項(xiàng)t,詞項(xiàng)權(quán)重計(jì)算為
,其中
是將隱藏狀態(tài)映射到浮點(diǎn)數(shù)的矩陣。如果詞項(xiàng)t在查詢中出現(xiàn)多次,只保留其最大權(quán)重。查詢和段落之間的相關(guān)性得分通過查詢和段落中共存詞項(xiàng)的聯(lián)合重要性計(jì)算:
。
- 多向量檢索:作為密集檢索的擴(kuò)展,多向量方法利用整個(gè)輸出嵌入來表示查詢和段落:
,其中
是可學(xué)習(xí)的投影矩陣。使用延遲交互計(jì)算細(xì)粒度的相關(guān)性得分:
,其中N和M分別是查詢和段落的長度。
自知識(shí)蒸餾
嵌入模型被訓(xùn)練為將正樣本與負(fù)樣本區(qū)分開,對(duì)于每種檢索方法,期望為查詢的正樣本分配比負(fù)樣本更高的得分。訓(xùn)練過程旨在最小化 InfoNCE 損失,其一般形式為:
其中,和
分別表示查詢q的正樣本和負(fù)樣本;s(?)是
中的任意一個(gè)函數(shù)。
不同檢索方法的訓(xùn)練目標(biāo)可能會(huì)相互沖突。為了促進(jìn)多種檢索功能的優(yōu)化,在訓(xùn)練Bge-M3時(shí),提出在自知識(shí)蒸餾的基礎(chǔ)上統(tǒng)一訓(xùn)練過程。具體步驟如下:
1. 集成預(yù)測(cè)得分:將不同檢索方法的預(yù)測(cè)得分整合為更準(zhǔn)確的相關(guān)性得分:
2. 計(jì)算損失函數(shù):計(jì)算的加權(quán)和作為沒有自知識(shí)蒸餾的損失:
3. 修改損失函數(shù):使用集成得分作為教師,修改每個(gè)檢索方法的損失函數(shù):
其中,p(?)是 softmax 激活函數(shù);是
中的任意一個(gè)成員。
4. 最終損失函數(shù):
- 進(jìn)一步整合并歸一化修改后的損失函數(shù):
- 最終損失函數(shù)為:
訓(xùn)練過程分為兩個(gè)階段:
- 預(yù)訓(xùn)練階段:使用大規(guī)模無監(jiān)督數(shù)據(jù)對(duì)文本編碼器(XLM-RoBERTa 模型,通過 RetroMAE 方法調(diào)整)進(jìn)行預(yù)訓(xùn)練,僅以對(duì)比學(xué)習(xí)的基本形式訓(xùn)練密集檢索。
- 微調(diào)階段:應(yīng)用自知識(shí)蒸餾,對(duì)嵌入模型進(jìn)行微調(diào),以建立三種檢索功能。在此階段使用標(biāo)注數(shù)據(jù)和合成數(shù)據(jù),并引入硬負(fù)樣本。
高效批處理
嵌入模型需要從多樣化和大規(guī)模的多語言數(shù)據(jù)中學(xué)習(xí),以充分捕捉不同語言的通用語義。為了確保文本嵌入的區(qū)分性,模型需要盡可能保持大批量(引入大量批內(nèi)負(fù)樣本)。然而,由于 GPU 內(nèi)存和計(jì)算能力的限制,傳統(tǒng)的做法是將輸入數(shù)據(jù)截?cái)酁槎绦蛄校詫?shí)現(xiàn)高訓(xùn)練吞吐量和大批量。這種方法雖然有效,但對(duì)于BGE-M3-Embedding 來說并不適用,因?yàn)樗枰瑫r(shí)處理短序列和長序列數(shù)據(jù),以有效支持不同粒度的輸入。
為了克服上述挑戰(zhàn),本文提出了一種高效的批處理策略,具體包括以下幾個(gè)關(guān)鍵步驟:
- 按序列長度分組:訓(xùn)練數(shù)據(jù)按序列長度分組,生成小批量時(shí)從同一組中采樣,顯著減少序列填充,提高了 GPU 的利用效率。
- 固定隨機(jī)種子:在為不同 GPU 采樣訓(xùn)練數(shù)據(jù)時(shí),固定隨機(jī)種子,確保負(fù)載平衡并減少每個(gè)訓(xùn)練步驟的等待時(shí)間。
- 子批量處理:在處理長序列訓(xùn)練數(shù)據(jù)時(shí),將小批量進(jìn)一步劃分為子批量,然后使用梯度檢查點(diǎn)(gradient checkpointing)技術(shù),迭代編碼每個(gè)子批量,顯著減少內(nèi)存占用的同時(shí)顯著增加了批量大小。
- 跨 GPU 廣播:在分布式訓(xùn)練環(huán)境中,將來自不同 GPU 的嵌入結(jié)果進(jìn)行廣播,使每個(gè) GPU 都可以獲取所有嵌入結(jié)果。通過廣播,每個(gè) GPU 上的模型可以訪問更多的負(fù)樣本,從而顯著擴(kuò)大了批內(nèi)負(fù)樣本的規(guī)模,有助于提高嵌入的區(qū)分性,進(jìn)一步提升模型的性能。
2、實(shí)驗(yàn)結(jié)果
多語言檢索
使用 MIRACL,包含 18 種語言的臨時(shí)檢索任務(wù),每個(gè)任務(wù)由同一語言的查詢和段落組成。
- Dense 方法:M3-Embedding 在密集檢索上表現(xiàn)出色,平均性能優(yōu)于所有基線方法。例如,在英語上與 E5mistral-7b 相比,M3-Embedding 在其他語言上表現(xiàn)更好。
- Sparse 方法:M3-Embedding 的稀疏檢索功能優(yōu)于傳統(tǒng)的 BM25 方法。
- Multi-vec 方法:多向量檢索進(jìn)一步提升了性能,通過細(xì)粒度的交互計(jì)算相關(guān)性得分。
- 混合方法:Dense+Sparse 和 All 方法通過結(jié)合不同方法的優(yōu)勢(shì),進(jìn)一步提升了檢索性能,其中 All 方法表現(xiàn)最佳。
跨語言檢索
使用 MKQA,包含 25 種非英語語言的查詢,目標(biāo)是從英文維基百科中檢索包含答案的段落。
M3-Embedding 在跨語言檢索任務(wù)中表現(xiàn)出色,尤其是在密集檢索和多向量檢索方面。結(jié)合不同方法的混合檢索進(jìn)一步提升了性能。此外,M3-Embedding 在低資源語言上的表現(xiàn)相對(duì)穩(wěn)定。
多語言長文檔檢索
BGE-M3可以支持長達(dá)8192的輸入文檔,同時(shí)長文檔檢索的效果要顯著優(yōu)于此前的基線模型。非常有意思的是,從實(shí)驗(yàn)結(jié)果可以觀察到,稀疏檢索(Sparse)的效果要顯著高于稠密檢索(Dense),這說明關(guān)鍵詞信息對(duì)于長文檔檢索極為重要。
消融實(shí)驗(yàn)
自知識(shí)蒸餾和多階段訓(xùn)練對(duì) M3-Embedding 的性能提升起到了關(guān)鍵作用。通過這些技術(shù),M3-Embedding 在多語言檢索、跨語言檢索和長文檔檢索任務(wù)中均表現(xiàn)出色。
3、總結(jié)
BGE-M3-Embedding模型的目的是構(gòu)建一個(gè)多功能、多語言、多粒度的文本嵌入模型,能夠高效地支持各種信息檢索和自然語言處理任務(wù)。它的意義在于:
- 多功能性:統(tǒng)一了多種檢索功能,適應(yīng)不同的檢索需求。
- 多語言支持:支持超過100種語言,適用于全球化的信息檢索和多語言對(duì)話系統(tǒng)。
- 多粒度處理:能夠處理從短句子到長文檔的不同粒度的輸入,適應(yīng)各種文本長度的需求。
- 實(shí)際應(yīng)用:在搜索引擎、問答系統(tǒng)、推薦系統(tǒng)、文檔分析和多語言對(duì)話系統(tǒng)等領(lǐng)域具有廣泛的應(yīng)用前景。
從評(píng)測(cè)結(jié)果來看,BGE-M3全面應(yīng)用于檢索任務(wù),包括訓(xùn)練數(shù)據(jù)也全是檢索的相關(guān)任務(wù),非常適合檢索增強(qiáng)任務(wù)。
在實(shí)際應(yīng)用中,BGE-M3也是text-embedding模型中下載量最高的模型,是目前最適合檢索增強(qiáng)任務(wù)的開源文本嵌入模型之一。