作者 | ethanntang

一、引言

1.概述

隨著大模型時代的爆發(fā)，檢索增強生成技術(shù)（RAG）在大語言模型中廣泛應用。RAG是一種性價比極高的方案，在大語言模型中占據(jù)重要地位。Embedding模型作為RAG中檢索召回的重要一環(huán)，扮演著極其關(guān)鍵的角色。更加準確的Embedding模型在抑制模型幻覺、增強新熱知識表現(xiàn)、提升封閉領域回答能力等方面都能發(fā)揮優(yōu)勢。

為了提升RAG系統(tǒng)的性能表現(xiàn)，我們近期針對如何訓練更強的Embedding模型進行探索，訓練得到了目前最強中文Embedding模型「Conan-Embedding」，該模型已在C-MTEB上達到SOTA。

圖1 C-MTEB榜單結(jié)果

模型鏈接：TencentBAC/Conan-embedding-v1 · Hugging Face （模型已上傳開源，歡迎試用）

評測榜單：MTEB Leaderboard - a Hugging Face Space by mteb

2.背景介紹

Embedding是一種將高維數(shù)據(jù)轉(zhuǎn)換為低維向量表示的技術(shù)，它在自然語言處理、計算機視覺等領域有廣泛應用。例如，在文本中，一個詞或短語會被轉(zhuǎn)換成一個固定長度的向量，這個向量能夠捕捉詞義和上下文信息。通過 Embedding 模型，可以計算句子間的相似度，應用于檢索，分類，召回，排序等諸多任務。

目前，Embedding模型已經(jīng)取得了顯著進展。例如，Word2Vec、GloVe 等模型能夠生成高質(zhì)量的詞嵌入向量。Transformer架構(gòu)的出現(xiàn)，進一步推動了Embedding模型的進步。BERT、GPT等預訓練語言模型通過大規(guī)模數(shù)據(jù)集訓練，生成了更加豐富和精確的詞嵌入表示。

通常，Embedding模型是通過對比學習來訓練的，而負樣本的質(zhì)量對模型性能至關(guān)重要。難負例挖掘就是利用 Teacher 模型，來找到與Query有一定相關(guān)性但不如正樣本相關(guān)的段落，從而使對比損失更難區(qū)分正例和負例。盡管難負例挖掘非常重要，但在Embedding工作中，這些方法往往被忽視，研究通常集中在模型架構(gòu)、微調(diào)方法和數(shù)據(jù)選擇上。

二、主要方法

圖2 訓練方法概覽

在弱監(jiān)督訓練階段，我們收集了 7.5 億對數(shù)據(jù)集，并從中挑選出 4 億對。在有監(jiān)督訓練階段，使用動態(tài)難負例挖掘策略來更精確地微調(diào)模型。

Language Embedding模型訓練通常采用多階段方案，分為弱監(jiān)督的預訓練以及有監(jiān)督的精調(diào)訓練。我們遵循這種訓練方式，將訓練分為預訓練和微調(diào)兩個階段。以下介紹了我們的詳細訓練流程和主要的方法，包含了動態(tài)的難負例挖掘和跨GPU的多任務均衡訓練方式。

1.訓練流程

(1) 預訓練

如圖2(a)所示，在預訓練階段，我們首先使用Internlm2.5中描述的標準數(shù)據(jù)過濾方法。首先，通過文檔提取和語言識別進行格式化處理；接著，在基于規(guī)則的階段，文本會經(jīng)過規(guī)范化和啟發(fā)式過濾；然后，通過MinHash方法進行去重；在安全過濾階段，執(zhí)行域名阻止、毒性分類和色情內(nèi)容分類；最后，在質(zhì)量過濾階段，文本會經(jīng)過廣告分類和流暢度分類，以確保輸出文本的高質(zhì)量。通過過濾，我們篩選了約 4.5 億對數(shù)據(jù)，留存率約60%。

bge-large-zh-v1.5 是由智源發(fā)布的廣泛使用的基礎embedding 模型。我們認為，該模型在對數(shù)據(jù)進行評分時，能夠有效地識別并保留高質(zhì)量數(shù)據(jù)。在數(shù)據(jù)經(jīng)過標準過濾后，我們使用bge-large-zh-v1.5模型對每一條數(shù)據(jù)進行評分，丟棄所有得分低于0.4的數(shù)據(jù)。通過評分，我們篩選了約 4 億對數(shù)據(jù)，留存率約 89%。

在預訓練階段，為了高效且充分地利用數(shù)據(jù)，我們使用InfoNCE Loss with In-Batch Negative：

其中是title，input，question，prompt等，是對應的 content，output，answer，response等，認為是正樣本；是同 batch 其他樣本的content，output，answer，response，認為是負樣本。In-Batch Negative InfoNCE Loss 是一種用于對比學習的損失函數(shù)，它利用了 mini-batch 中的其他樣本作為負樣本來優(yōu)化模型。具體來說，在每個 mini-batch 中，除了目標樣本的正樣本對外，其余樣本都被視為負樣本。通過最大化正樣本對的相似度并最小化負樣本對的相似度，In-Batch Negative InfoNCE Loss 能夠有效地提高模型的判別能力和表征學習效果。這種方法通過充分利用 mini-batch 中的樣本，提升了訓練效率并減少了對額外負樣本生成的需求。

(2) 有監(jiān)督精調(diào)

微調(diào)階段，我們針對不同的下游任務進行特定任務的微調(diào)。如圖2(b)所示，我們參考以往的工作，并在其基礎上移除了分類（CLS）任務，將任務分為兩類：檢索（Retrieval）和語義文本相似性（STS）。檢索任務包括查詢、正樣本和負樣本，經(jīng)典的損失函數(shù)是InfoNCE Loss。STS任務涉及區(qū)分兩段文本之間的相似性，經(jīng)典的損失函數(shù)是交叉熵損失。在STS任務上，根據(jù)以往工作的結(jié)論，CoSENT損失略優(yōu)于交叉熵損失。因此，我們也采用CoSENT損失來優(yōu)化STS任務：

其中是縮放溫度，是余弦相似度函數(shù)，是 和之間的相似性。

2.動態(tài)難負例挖掘訓練

在embedding模型準備數(shù)據(jù)時，難負例挖掘用于為query選擇負樣本。其思想是使用一個 teacher模型來找到與query有一定相關(guān)性但不如正樣本相關(guān)的段落，從而使對比損失更難區(qū)分正例和負例。這些難負例應該比隨機負例更難與正例區(qū)分，從而帶來更高效和更有效的微調(diào)。

先前的工作基本都在數(shù)據(jù)預處理階段進行難負例挖掘。對于Embedding模型來說，一個既定的權(quán)重狀態(tài)下的模型，難負例是一定的。然而，隨著訓練的進行，每當模型權(quán)重更新時，當前權(quán)重下的模型對應的難負例就會變化。在數(shù)據(jù)預處理階段挖掘的難負例在經(jīng)過訓練迭代后，就會變得不那么難了?；谶@一觀點，我們提出了一種動態(tài)難負例挖掘方法。對于每個數(shù)據(jù)，我們記錄當前難先前的工作基本都在數(shù)據(jù)預處理階段進行難負例挖掘。對于Embedding模型來說，一個既定的權(quán)重狀態(tài)下的模型，難負例是一定的。然而，隨著訓練的進行，每當模型權(quán)重更新時，當前權(quán)重下的模型對應的難負例就會變化。在數(shù)據(jù)預處理階段挖掘的難負例在經(jīng)過訓練迭代后，就會變得不那么難了。

基于這一觀點，我們提出了一種動態(tài)難負例挖掘方法。對于每個數(shù)據(jù)，我們記錄當前難負例與Query的平均分數(shù)。每 100 次迭代后，如果分數(shù)的 1.15 倍小于初始分數(shù)且分數(shù)絕對值小于 0.8時，我們就認為該負例不再困難，并進行新一輪的難負例挖掘。每次進行動態(tài)難負例挖掘時，如果需要替換難負例，我們使用到個案例作為負樣本，其中表示第次替換，表示每次使用的難負例數(shù)。整個過程產(chǎn)生的成本僅相當于一個 step迭代。相比于In-Batch Negative InfoNCE Loss，我們認為更高質(zhì)量（更符合當前模型權(quán)重下)的難負例更為重要。圖3 展示了動態(tài)難負例挖掘與標準難負例挖掘的樣本正負例Score - Steps 變化曲線。可以看到，隨著步數(shù)的增加，Standard-HNM的負例評分不再下降，而是出現(xiàn)震蕩，這表明模型對該批負例的學習已完成。而Dynamic-HDM在檢測到負例學習完畢后，會進行難負例的替換。。

圖3 展示了動態(tài)難負例挖掘與標準難負例挖掘的樣本正負例Score - Steps 變化曲線?？梢钥吹?，隨著步數(shù)的增加，Standard-HNM的負例評分不再下降，而是出現(xiàn)震蕩，這表明模型對該批負例的學習已完成。而Dynamic-HDM在檢測到負例學習完畢后，會進行難負例的替換。

圖3 動態(tài)難負例挖掘與標準難負例挖掘的樣本正負例Score-Steps變化曲線。

在訓練過程中，每 100steps 檢查一次難負例。當分數(shù)的 1.15 倍小于初始得分，且分數(shù)絕對值小于 0.8 時，我們認為該負例不再困難，并替換成新的難負例。

3.跨GPU的Batch均衡訓練

圖4 跨 GPU 的Batch均衡訓練的示意圖

對于Retri任務，我們利用多個 GPU 來合并更多難負例。對于 STS 任務，我們增加BatchSize以包含更多樣本進行學習，以獲得更魯棒的排名關(guān)系。

為了更好地利用難樣本，我們采用了跨 GPU 批次平衡損失 (CBB)。之前的方案通常在訓練流程中隨機的在每個Batch中分配一個任務。例如：在iter0中采樣STS的樣本，并使用STS對應Loss進行反向傳播獲取梯度并更新權(quán)重，而iter1中分配了Retri任務或者CLS任務，我們稱之為順序隨機任務訓練。這樣訓練幾乎一定會導致單次的優(yōu)化搜索空間與Embedding模型的全局優(yōu)化搜索空間不一致，從而導致訓練過程的震蕩以及無法求得全局最優(yōu)解。我們在之后的分析中展現(xiàn)了這一現(xiàn)象。

為此，我們考慮在每次的Forward-Loss-Backward-Update的更新過程中都均衡的引入每一個任務，以此來獲得穩(wěn)定的搜索空間，并盡可能的縮小單次模型更新方向和全局最優(yōu)解的一致性。因此，CBB策略不僅考慮了不同 GPU 之間的通信，還考慮了不同任務之間的通信，從而實現(xiàn)了更好的Batch均衡。如圖4所示，為了在檢索任務中利用更多難樣本，我們確保 GPU（gpu0、gpu1、gpu2、gpu3）各自具有不同的負樣本，同時共享相同的查詢和相同的正樣本。對于Retri任務，每個 GPU 計算對應Batch的Loss，然后將結(jié)果匯總到 gpu1 上。對于STS任務，在gpu4上，運行STS任務并獲得對應Loss。最終匯總并計算當前Iter的合并 CBB Loss。對應公式如下：

其中是 Query 和正樣本段落之間的評分函數(shù)，通常定義為余弦相似度，是查詢和正樣本段落共享的GPU數(shù)量，是縮放溫度。根據(jù)經(jīng)驗，我們將設置為 0.8。

三、實驗

1.實驗細節(jié)

與大多數(shù)Embedding模型一樣，Conan-Embedding也采用BERT模型作為基礎模型，并使用FC Layer將輸出維度從1024擴展到1792。模型的參數(shù)量為326M。Conan-Embedding的最大輸入長度為 512 個 token。此外，受到 OpenAI的text-embedding-v3的啟發(fā)，我們還利用了多尺度表征學習（Matryoshka Representation Learning, MRL）技術(shù)來實現(xiàn)靈活的輸出維度長度，提升模型表征性能和魯棒性。對于 MRL 訓練，表示維度配置為256、 512、 768、1024、1536 和 1792。

為了提高效率，我們使用了混合精度訓練和 DeepSpeed ZERO-Stage 1。

弱監(jiān)督預訓練階段，我們使用 AdamW優(yōu)化器，初始學習率為 1e-5，Warmup設置為 0.05，Decay設置為 0.001。整個預訓練過程使用了 64張華為Ascend 910B GPU，單次精調(diào)訓練約消耗138 個小時。

有監(jiān)督精調(diào)階段，檢索任務的BatchSize設置為 4，STS 任務的BatchSize設置為32。我們使用了與預訓練階段相同的優(yōu)化器參數(shù)和學習率。整個微調(diào)過程使用了 16 張華為Ascend 910B GPU，單次精調(diào)訓練約消耗13 個小時。

2.數(shù)據(jù)情況

在預訓練階段，我們收集了約7.5 億個文本對，包含Wudao、Zhihu-KOL、SimCLUE等，數(shù)據(jù)分為標題-內(nèi)容對、輸入-輸出對和問答對等。我們還發(fā)現(xiàn)，高質(zhì)量的 LLM 指令調(diào)優(yōu)數(shù)據(jù)（例如：提示-響應對）經(jīng)過規(guī)則過濾和篩選后，可以提升 Embedding 模型的性能。此外，我們利用現(xiàn)有的文本語料庫，使用 LLM 生成了一批數(shù)據(jù)。詳細的數(shù)據(jù)描述可以在下表中展示：

表1 預訓練數(shù)據(jù)概覽

在精調(diào)階段，為了讓模型更適應各種任務，我們選擇了常見的Retri、CLS和 STS數(shù)據(jù)集。對于CLS任務，我們將它與Retri合并，將同一類別的數(shù)據(jù)視為文本正例，將不同類別的數(shù)據(jù)視為文本負例。微調(diào)階段使用的數(shù)據(jù)量如下表所示：

表2 精調(diào)訓練數(shù)據(jù)概覽

3.結(jié)果

(1) 消融實驗結(jié)果

為了證明我們方法的有效性，我們在 CMTEB 基準上進行了全面的消融研究。如下表所示，動態(tài)難負樣本挖掘（DHNM）和跨 GPU 批次平衡（CBB）都明顯優(yōu)于直接微調(diào)模型的Vanilla方法。除此之外，Conan-Embedding在Retri和Rerank任務中表現(xiàn)出顯著的改進，這表明隨著負樣本數(shù)量的增加和質(zhì)量的提高，模型能夠看到更具有學習價值的難負樣本，從而增強了其召回能力。

表3 消融實驗結(jié)果

Baseline 表示我們在預訓練階段之后的結(jié)果；Vanilla 表示直接使用標準對比和 CoSENT 損失進行微調(diào)；

DHNM 表示僅使用動態(tài)硬負挖掘方法；CBB Loss 表示僅使用跨 GPU 批次平衡損失。

(2) C-MTEB結(jié)果

MTEB（Massive Text Embedding Benchmark）是評估大規(guī)模文本Embedding模型的最權(quán)威和最流行的基準。MTEB創(chuàng)建了一個中文 Embedding評估Benchmark，稱為 C-MTEB。C-MTEB 有 35 個數(shù)據(jù)集，涵蓋 6 個類別：分類、聚類、文本對分類、重排序、檢索和 STS。下表展示了我們的模型與 C-MTEB 基準上其他模型的比較。我們的模型在幾乎所有任務上都超越了之前最先進的模型，包含參數(shù)量更大的Decoder-Only模型。

表4 C-MTEB結(jié)果

4.分析

為了更好地評估跨GPU多任務批次均衡的效果，我們在圖5展示了使用CBB策略前后的Loss-Iter曲線。loss_retri和loss_STS代表在不啟用CBB策略時，多個任務隨機按Iter進行訓練時，兩個任務的獨立Loss。可以觀察到，在不啟用CBB策略時，每個獨立任務的Loss震蕩都比較嚴重，并且下降緩慢，且不同步。以上現(xiàn)象說明，不同的任務之間優(yōu)化目標存在GAP，優(yōu)化方向不一致，且每個單一優(yōu)化目標和全局優(yōu)化目標不一致。因此，順序隨機任務訓練不能在優(yōu)化中取得近似的全局最優(yōu)結(jié)果。loss_cross代表啟用CBB策略時訓練過程中的Loss-Iter曲線，可以看到，隨著訓練的進行，Loss在幾乎在非常平穩(wěn)地持續(xù)的下降，最終的Loss（0.08）遠小于Retri和STS Loss之和（0.38）。CBB策略可以看作是一種正則化策略。

圖5 使用跨GPU多任務批次均衡（CBB）方法前后的損失曲線比較

四、展望

在實際業(yè)務中，由于業(yè)務體系龐大、業(yè)務標準靈活多變、不同業(yè)務標準不一，運用Embedding模型配合RAG技術(shù)來提升實際業(yè)務的召回性能、快速適應標準變化、降低模型訓練開銷，是優(yōu)化成本與精度的不錯選擇。

另外，Embedding模型也可以在多業(yè)務場景下獨立提供召回服務，目前我們已有Conan- Embedding推理服務上線提供給業(yè)務方使用，并能適配多種特征長度的需求。我們正在推進Conan-Embedding在實際業(yè)務中的落地，例如：不同體系標準的標簽業(yè)務、推薦系統(tǒng)相似樣本召回、以及通用RAG系統(tǒng)等等。

最后，歡迎大家聯(lián)系我們使用Conan-Embedding，并提出寶貴意見。當前訓練數(shù)據(jù)可能仍存在短板，不能直接在所有類型業(yè)務上都取得最佳表現(xiàn)，其中還面臨一些有待解決問題和算法挑戰(zhàn)，非常歡迎大家與我們進行合作，推進Conan-Embedding技術(shù)探索和迭代，共同解決業(yè)務落地中的各種問題。