在一天內構建特定領域的嵌入模型

在一天內構建特定領域的嵌入模型

只需單個 GPU 和不到一天的訓練時間，您就可以將通用嵌入模型轉換為真正理解您專業領域的模型，且無需手動標記。為了幫助您快速上手，我們還發布了一個現成的合成訓練數據集，該數據集是使用此流程從 NVIDIA 的公開文檔中生成的。使用這些數據和方案，我們觀察到 Recall@10 和 NDCG@10 均有超過 10% 的提升。Atlassian 應用此方案在其 JIRA 數據集上進行微調，將 Recall@60 從 0.751 提高到 0.951，提升了 26% —— 且僅使用了單個 GPU。

🔗 數據集與代碼快速連結：

🧑‍💻 方案整合的開源項目：

📋 前提條件：

閱讀完本文後，您將了解如何：📄 從領域文檔生成訓練數據（無需標籤數據）🎯 使用難負例挖掘（Hard Negative Mining）進行有效的對比訓練🔗 通過多跳查詢（Multi-hop queries）提高嵌入質量⚙️ 微調雙編碼器（Bi-encoder）嵌入模型📊 評估微調是否改善了檢索效果🚀 在您的流水線中部署微調後的模型

⚙️ 設置

在本教程中，我們將微調基礎模型 Llama-Nemotron-Embed-1B-v2 —— 這是一個擁有 10 億參數的嵌入模型，平衡了質量和推理成本。要開始使用，請遵循此設置指南。

📚 步驟 1：從文檔生成訓練數據

微調嵌入模型需要數千個（查詢，相關文檔）對。大多數使用場景並沒有現成的這類數據。手動創建既昂貴、緩慢，且往往會受到標註者對「相關性」個人理解的偏見影響。與其手動標記數據，您可以使用 LLM (nvidia/nemotron-3-nano-30b-a3b) 閱讀您的文檔並自動生成高質量的合成問答對。

它是如何運作的？

在後台，這運行了一個由 NeMo Data Designer 驅動的四階段合成數據生成 (SDG) 流水線：

輸出結果長什麼樣？

源文檔片段：

H100 GPU 在 SXM 封裝形式下的熱設計功耗 (TDP) 為 700W。冷卻解決方案必須在持續工作負載下將結溫維持在 83°C 以下。對於每節點超過 4 個 GPU 的密集部署，建議使用液冷，因為風冷無法在標準 2U 機箱配置中散發足夠的熱量。

生成的問答對：

注意區別：第一個問題是簡單的事實查找。第二個問題則需要多跳、因果推理。該流水線會生成這兩種類型，並具有可配置的複雜度等級 (2–5) 和跳數 (1–3)。每個問答對隨後都會經過質量評估，接收關於相關性、準確性、上下文支持和清晰度的子評分，以及一個總分。只有達到閾值的問答對才會被納入訓練。

⛏️ 步驟 2：挖掘難負例（以及為什麼它們很重要）

如果您僅使用正例對（查詢 + 正確文檔）訓練嵌入模型，它會學會區分明顯不同的文檔，但在困難案例上會失敗 —— 即那些看起來相關但並非正確答案的段落。在真實的檢索系統中，這些「差一點就對」的文檔正是導致錯誤答案的原因。難負例挖掘會找出這些令人困惑的段落，以便模型學習如何區分它們。

上述命令會自動運行三個子步驟：

2a. 訓練 / 驗證 / 測試集分割

生成的問答對被分為訓練集 (80%) 和測試集 (20%)。測試集被格式化為兼容 BEIR 的基準測試，以便在步驟 5 中進行標準化評估。

2b. 難負例挖掘

使用基礎嵌入模型，流水線會：

結果：難負例是與正例最相似、但仍安全低於正例得分上限的非正例段落。它們是當前模型認為高度相關但並非標記答案的段落。

為什麼這有效：在簡單負例（完全不相關的段落）上訓練無法讓模型學到新知識。在難負例上訓練會迫使它學習領域中重要的細微差別。例如，在醫學語料庫中，關於「二型糖尿病的二甲雙胍劑量」的問題可能會涉及關於「二甲雙胍副作用」或「一型糖尿病的胰島素劑量」的難負例 —— 它們很接近但有本質區別。95% 的邊際上限可防止挖掘程序選擇與正例過於接近的段落，因為那些段落可能是正確答案，只是在 SDG 期間未被標記。

2c. 多跳展開 (Multi-Hop Unrolling)

多跳問題引用多個正例文檔。例如，「第 3.2 節中的熱管理系統與第 5.1 節中描述的功率限制有何關係？」這類問題有兩個正例段落。

展開會為每個（查詢，正例文檔）對創建一個訓練示例，使對比損失函數能獨立看待每個正例。一個具有 2 個正例文檔的問題會變成 2 個訓練示例，每個示例具有相同的難負例但正例不同。

最終輸出是一個可用於訓練的 JSON 文件：

🔍 步驟 3：理解多跳問題以及為何它們能改善檢索

標準的嵌入微調為每個段落生成一個問題並訓練模型進行匹配。這適用於簡單的事實查找，但真實用戶會提出跨越多個文檔或章節的複雜問題。如果模型只見過單跳訓練數據，它將難以檢索到這些複雜查詢的所有相關段落。

SDG 流水線默認生成 1 到 3 跳的問題：

每一跳都由其自身的上下文摘要和段落 ID 進行追蹤，因此訓練數據保留了完整的推理鏈。在展開（步驟 2c）之後，每個（問題，相關段落）對都成為一個獨立的訓練信號，教導模型所有這些段落都與該多跳查詢相關。

微調後的模型學會檢索語義相關的文檔，而不僅僅是詞彙相似的文檔。

🧠 步驟 4：微調嵌入模型

對比學習的工作原理

訓練使用帶有對比損失的雙編碼器架構。

0.02 的溫度設定是刻意激進的，它會產生非常尖銳的概率分佈。這之所以有效，是因為步驟 2 中的難負例質量很高：它們是真正令人困惑的段落，模型需要強大的梯度才能學會區分。

關鍵超參數

小數據集的自動縮放

如果您的數據集少於 2,000 個訓練示例，流水線會自動：

這意味著您可以從一個小型語料庫（50–100 個文檔）開始進行快速概念驗證，稍後再進行擴展。

📈 步驟 5：衡量改進效果

微調真的有幫助嗎？讓我們通過在預留的測試集上運行標準化評估，比較基礎模型與微調後的檢查點：

評估使用 BEIR 框架，並計算 k = 1, 5, 10 和 100 時的四項標準信息檢索指標：

成功的微調通常會在不到 1 天內使 nDCG@10 和 Recall@10 提升 15%。

使用 Retrieval Synthetic NVDocs 的結果：

如果數據沒有改善怎麼辦？

該流水線易於迭代：

🏆 現實世界案例：Atlassian

此方案已由 Atlassian 在真實企業數據上得到驗證。他們按照上述相同階段，使用單個 NVIDIA A100 80GB GPU，將此流水線應用於在公開的 Jira 數據集上微調 Llama-Nemotron-Embed-1B-v2。

Recall@60 從 0.751 躍升至 0.951 —— 增長了 26.7%。微調後的模型在 95.1% 的查詢中，能在前 60 個結果內檢索到正確文檔，而基礎模型僅為 75.1%。對於支撐 Jira 搜索的檢索系統來說，這直接轉化為為數百萬用戶提供更相關的結果。在他們的博客文章《為數百萬 Rovo 用戶推進語義搜索》中可以找到更多細節。

🚀 步驟 6：導出與部署

PyTorch 檢查點對於評估很好，但對於生產環境來說太慢了。最後兩個階段將轉換模型並將其置於 API 之後提供服務。

導出至 ONNX / TensorRT

這會將微調後的檢查點導出為 ONNX (opset 17)。可選地，它會編譯一個 TensorRT 引擎以實現最大推理吞吐量，並具有針對批次大小 (1–64) 和序列長度 (3–256) 的可配置優化配置：

使用 NVIDIA NIM 部署

導出的模型部署在 NVIDIA NIM 容器中 —— 這是一個生產級的推理微服務，公開了與 OpenAI 兼容的 /v1/embeddings 端點：

運行後，任何客戶端都可以調用它：

由於 NIM 提供與 OpenAI 兼容的 API，您可以將其放入任何使用嵌入 API 格式的現有 RAG 流水線中 —— 無需更改代碼。

驗證部署準確性

流水線包含一個 NIM 準確性驗證步驟，該步驟會針對部署的端點運行相同的 BEIR 評估：

這可以捕捉到 ONNX/TensorRT 轉換過程中產生的任何準確性損失。在容差範圍內匹配的指標（@1 為 0.03，@5+ 為 0.01）會被標記為勾選；超出轉換噪聲的偏差則會被標記。

總結

完整的嵌入微調流水線可以通過六個命令運行，從原始文檔到部署模型。

預期時間和資源

總計：不到一天，大部分時間是無需人工干預的訓練。對於小型語料庫（約 500 個文檔），整個流水線大約在 2-3 小時內完成。

該流水線可以端到端運行，但每個階段也可以根據您的起點獨立執行。例如，如果您有原始文檔，可以從合成數據生成 (SDG) 開始；而已經包含難負例的數據集可以跳過前面的步驟直接進行微調。由於每個階段都使用 JSON、BEIR 和 ONNX 等標準格式，因此很容易集成自定義組件或在其他工作流中重用中間輸出。該方案在運行方式上也很靈活，支持在本地機器、Docker 容器或基於 Slurm 的集群上執行。

親自嘗試

如果您有領域文檔和一些時間，今天就可以生成您的第一批合成訓練數據！從文檔到部署的領域適應嵌入模型的完整流水線，在單個 GPU 上運行不到一天即可完成。您可以從我們準備好的 nvidia/Retrieval-Synthetic-NVDocs-v1 數據集開始，立即嘗試該流水線。讓我們知道您構建了什麼。如果您覺得有用，請為 Nemotron、NeMo Data Designer 和 NeMo Automodel 的倉庫點亮星星。

社群

· 註冊或登錄以發表評論