跳到正文
原文
Hugging Face Blog·· 2026-03-21精选AI 评分63

如何在一天内构建领域专用嵌入模型

Build a Domain-Specific Embedding Model in Under a Day

AI 导读

NVIDIA 发布一套在单张 GPU 上不到一天完成领域嵌入模型微调的流程,基于 Llama-Nemotron-Embed-1B-v2,无需人工标注。流程用 nemotron-3-nano-30b-a3b 从领域文档生成合成问答数据、挖掘难负样本并做多跳展开,再用对比学习微调,实测 nDCG@10 与 Recall@10 均提升约 10%。

推荐理由

NVIDIA 给出了从合成数据、难负样本挖掘到 ONNX/TensorRT 部署的完整嵌入模型微调配方,附真实评测数字。

正文 · AI 翻译

如果你正在构建一个 RAG(检索增强生成)系统,你很可能已经撞上了这堵墙:一切都能跑通……直到它不行。通用嵌入模型被训练来理解互联网,而不是你的合同、制造日志、专有化学配方或内部分类体系。它们能捕捉宽泛的语义相似度,却无法理解在你所在领域中至关重要的细粒度差异。当现成的模型无法有效捕捉领域特定细微差别时,微调嵌入模型可以提升检索管线的性能。尽管嵌入对 RAG 性能如此关键,整个流程却仍然出奇地碎片化,所需技能相当专业,时间投入也令人生畏。

只需一块 GPU 和不到一天的训练时间,你就能把一个通用嵌入模型转变为真正理解你所在领域的模型,而且无需人工标注。为了帮你快速上手,我们还发布了一个开箱即用的合成训练数据集,它是使用这套完全相同的流程从 NVIDIA 公开文档生成的。使用这份数据和这套方法,我们在 Recall@10 和 NDCG@10 上都看到了超过 10% 的提升。Atlassian 将这套方法应用于他们的JIRA 数据集进行微调,将 Recall@60 从 0.751 提升到 0.951,提升了 26%——而且只用了单块 GPU。

🔗数据集和代码快速链接:

  1. 嵌入模型
  2. GitHub
  3. 基于 NVIDIA 公开文档的合成数据集

🧑‍💻开源项目方法集成了:

  1. NeMo Data Designer 用于合成数据生成
  2. NeMo Automodel 用于嵌入模型训练
  3. BEIR 用于信息检索评估
  4. NeMo Export-Deploy 用于 ONNX/TensorRT 转换
  5. NVIDIA NIM 用于生产推理服务

📋前提条件:

  • 一个包含领域文档的目录(文本文件 - .txt、.md 或类似格式)
  • 一个有效的 NVIDIA API 密钥(可在 build.nvidia.com 免费获取)
  • NVIDIA Ampere GPU or newer with at least 80GB memory (with Compute Capability >= 8.0)
    • 本教程已在 1xA100(80GB)和 1xH100(80GB)上测试通过

读完本文后,你将学会如何:
📄 在没有标注数据的情况下从领域文档生成训练数据
🎯 使用难负例挖掘进行有效的对比训练
🔗 通过多跳查询提升嵌入质量
⚙️ 微调双编码器嵌入模型
📊 评估微调是否改善了检索
🚀 在你的管线中部署微调后的模型

⚙️环境设置

在本教程中,我们将微调基础模型 Llama-Nemotron-Embed-1B-v2 —— 一个在质量和推理成本之间取得平衡的 10 亿参数嵌入模型。要开始使用,请遵循这份设置指南。

📚 第 1 步:从文档生成训练数据

微调嵌入模型需要数千个(查询,相关文档)对。大多数用例并没有现成可用的这类数据。手动创建这类数据成本高昂、速度缓慢,而且常常会受到标注者个人对“相关”理解的影响而带有偏差。
与其手工标注数据,你可以使用 LLM(nvidia/nemotron-3-nano-30b-a3b)来阅读你的文档,并自动生成高质量的合成问答对。

nemotron embed sdg -c default corpus_dir=./data/my_domain_docs

它是如何工作的?

在后台,这运行的是一个由 NeMo Data Designer 驱动的四阶段合成数据生成(SDG)流水线:

sdg_pipeline

输出是什么样的?

源文档片段:

H100 GPU 的热设计功耗(TDP)在 SXM 形态下为 700W。散热方案必须确保在持续负载下结温保持在 83°C 以下。对于每节点超过 4 块 GPU 的高密度部署,建议采用液冷,因为在标准 2U 机箱配置中,风冷无法消散足够的热量。

生成的问答对:

{
  "question": "What cooling approach is recommended when deploying more than 4 H100 GPUs per server node?",
  "answer": "Liquid cooling is recommended for dense deployments exceeding 4 GPUs per node, as air cooling cannot dissipate sufficient heat in standard 2U chassis configurations.",
  "query_type": "contextual",
  "reasoning_type": "factual",
  "question_complexity": 3,
  "segment_ids": [1],
  "quality_score": 8.5
}
{
  "question": "How does the 700W TDP of the H100 SXM constrain the choice between air and liquid cooling in multi-GPU configurations?",
  "answer": "The 700W TDP generates substantial heat that must be dissipated to keep junction temperatures below 83°C. In dense configurations exceeding 4 GPUs per node, air cooling in standard 2U chassis cannot handle this thermal load, making liquid cooling necessary.",
  "query_type": "multi_hop",
  "reasoning_type": "causal",
  "question_complexity": 4,
  "segment_ids": [1, 2],
  "hop_count": 2,
  "quality_score": 9.0
}

注意其中的区别:第一个问题是简单的事实查找。第二个问题则需要多跳的因果推理。该流水线会同时生成这两种类型,并支持可配置的复杂度级别(2–5)和跳数(1–3)。随后每个问答对都会经过质量评估,获得相关性、准确性、上下文支持度和清晰度等子分数,以及一个总体分数。只有达到阈值的问答对才会被纳入训练。

⛏️ 第 2 步:挖掘难负样本(以及它们为何重要)

如果只用正样本对(查询 + 正确文档)来训练嵌入模型,它学会的是区分明显不同的文档,但在困难情况下会失败——那些看起来相关但并非正确答案的段落。在真实的检索系统中,这些近似但不正确的文档正是导致糟糕答案的根源。难负样本挖掘会找出这些容易混淆的段落,让模型学会区分它们。

nemotron embed prep -c default

上述命令会自动运行三个子步骤:

2a. 训练 / 验证 / 测试划分

生成的问答对被划分为训练集(80%)和测试集(20%)。测试集被格式化为与 BEIR 兼容的基准,以便在第 5 步中进行标准化评估。

2b. 难负样本挖掘

使用基础嵌入模型,该流水线会:

  1. 对语料库中的每个查询和每个段落进行嵌入。
  2. 计算每个查询与所有段落之间的相似度。
  3. 屏蔽掉每个查询已标注的正样本文档。
  4. 应用边界过滤:任何非正样本文档,如果其得分高于最小正样本得分的 95%,都会被剔除。这个排除区可以防止假负样本——即那些与正样本非常接近、可能实际上相关的未标注段落。
  5. 从剩余的候选项中,选取得分最高的 top-k 文档作为难负样本(默认每个查询 5 个)。

结果是:难负样本是最相似的非正样本段落,同时仍安全地低于正样本得分上限。它们是当前模型认为高度相关、但并非标注答案的段落。

为什么这有效:用简单负样本(完全不相关的段落)训练无法让模型学到新东西。用难负样本训练则迫使模型学习在你的领域中至关重要的细微差别。例如,在医学语料库中,关于“二甲双胍治疗 2 型糖尿病的剂量”的问题,其难负样本可能是关于“二甲双胍副作用”或“胰岛素治疗 1 型糖尿病的剂量”的内容——相近但存在关键差异。95% 的边界上限可以防止挖掘器选择与正样本过于接近的段落,因为这些段落实际上可能是正确答案,只是在 SDG 过程中未被标注。

2c. 多跳展开

多跳问题会引用多个正样本文档。例如,像“第3.2节中的热管理系统与第5.1节中描述的功率限制有何关联?”这样的问题就有两个正样本段落。

展开会为每个(查询,正样本文档)对生成一个训练示例,因此对比损失会独立地看待每个正样本。一个带有2个正样本文档的问题会变成2个训练示例,每个示例拥有相同的难负样本但不同的正样本。

最终输出是一个可直接用于训练的JSON文件:

{
  "question_id": "q42_0",
  "question": "How does the thermal management system in Section 3.2 relate to the power constraints described in Section 5.1?",
  "pos_doc": [{"id": "Section 3.2"}],
  "neg_doc": [{"id": "d_x7y8z9"}, {"id": "d_m4n5o6"}, {"id": "d_p1q2r3"}, {"id": "d_s4t5u6"}, {"id": "d_v7w8x9"}]
}, 
{
  "question_id": "q42_1",
  "question": "How does the thermal management system in Section 3.2 relate to the power constraints described in Section 5.1?",
  "pos_doc": [{"id": "Section 5.1"}],
  "neg_doc": [{"id": "d_x7y8z9"}, {"id": "d_m4n5o6"}, {"id": "d_p1q2r3"}, {"id": "d_s4t5u6"}, {"id": "d_v7w8x9"}]
}

🔍 第3步:理解多跳问题及其为何能提升检索效果

标准的嵌入微调会为每个段落生成一个问题,并训练模型将它们匹配起来。这对于简单的事实查找是有效的,但真实用户会提出跨越多个文档或章节的复杂问题。如果模型只见过单跳训练数据,它在为这些复杂查询检索所有相关段落时会很吃力。

SDG流水线默认生成1到3跳的问题:

  • 1跳:“H100 SXM的TDP是多少?”——由单个段落即可回答。
  • 2跳:“H100的TDP与高密度部署中的冷却需求有何关联?”——需要连接两个段落中的信息。
  • 3跳:“考虑到TDP、冷却约束和机架密度限制,标准数据中心单排最多可部署多少块H100 GPU?”——综合三个段落的信息。

每一跳都通过其各自的上下文摘要和段落ID进行追踪,因此训练数据保留了完整的推理链。在展开(第2c步)之后,每个(问题,相关段落)对都会成为一个独立的训练信号,教会模型这些段落全部都与该多跳查询相关。

微调后的模型学会检索上下文相关的文档,而不仅仅是词汇上相似的文档。

🧠 第4步:微调嵌入模型

nemotron embed finetune -c default

对比学习的工作原理

训练使用带有对比损失的双编码器架构。

contrastive_learning

0.02的温度值是刻意设置得比较激进的,它会产生非常锐利的概率分布。这样效果很好,因为第2步得到的难负样本质量很高:它们确实是容易混淆的段落,模型需要较强的梯度才能学会区分它们。

关键超参数

参数 默认值 说明
轮数 3 默认值针对小型示例数据集进行了调优;对于真实世界的数据,1–2轮通常就足够了,以避免过拟合。
学习率 1e-5 调优:可以尝试默认值的两倍和一半
学习率预热步数 5 设置为微调总步数的5-10%,以获得更好的早期训练稳定性
全局批量大小 128
每个查询的段落数 5 1个正样本 + 4个难负样本

检查点保存频率

如果配置中省略了ckpt_every_steps,检查点保存频率会自动设置:

  • 映射式数据集(已知长度):默认每个轮次保存一次。
  • 可迭代数据集(未知长度):默认训练期间保存两次。

这意味着你可以先用一个小型语料库(50–100个文档)快速进行概念验证,之后再扩大规模,而无需手动调整检查点设置。

📈 第5步:衡量改进效果

微调真的有帮助吗?让我们在留出测试集上运行标准化评估,对比基础模型与微调后的检查点,来一探究竟:

nemotron embed eval -c default

评估使用 BEIR 框架,并在 k = 1、5、10 和 100 时计算四个标准信息检索指标:

  1. nDCG@k:排序质量——最佳文档是否排在最前面?
  2. Recall@k:覆盖率——相关文档中有多大比例出现在前 k 个结果中?
  3. Precision@k:准确率——前 k 个结果中实际相关的比例是多少?
  4. MAP@k:所有查询的平均精度

成功的微调通常能在 <1 天内使 nDCG@10 和 Recall@10 提升 10%。

使用 Retrieval Synthetic NVDocs 的结果:

📊 Comparison (Base -> Fine-tuned)
============================================================
  NDCG:
    NDCG@1: 0.55178 → 0.60796 (+0.05618, +10.2%)
    NDCG@5: 0.51894 → 0.57689 (+0.05795, +11.2%)
    NDCG@10: 0.55506 → 0.61559 (+0.06053, +10.9%)
    NDCG@100: 0.60617 → 0.66567 (+0.05950, +9.8%)
  Recall:
    Recall@1: 0.28478 → 0.31547 (+0.03069, +10.8%)
    Recall@5: 0.54486 → 0.60288 (+0.05802, +10.6%)
    Recall@10: 0.62979 → 0.69296 (+0.06317, +10.0%)
    Recall@100: 0.81421 → 0.87020 (+0.05599, +6.9%)

如果数字没有提升怎么办?

该流程让迭代变得简单:

  • SDG 中质量分数低? 检查你的文档质量——干净、格式良好的文本能生成更好的合成数据。尝试更大、更强大的 LLM。
  • 训练数据不够? 向语料库添加更多文档,并重新运行 Stage 0。
  • 过拟合? 默认的 3 个 epoch 是针对小型示例数据集校准的;对于大多数真实世界数据,1–2 个 epoch 就足够了。也可以考虑提高质量阈值,只保留最好的训练样本。
  • 学习率不对? 试试将默认学习率翻倍或减半。

🏆 真实世界成果: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 搜索的检索系统而言,这直接为数百万用户带来了更相关的结果。更多细节请见他们的博客文章Advancing semantic search for millions of Rovo users。

ATL result

🚀 第 6 步:导出与部署

PyTorch 检查点很适合评估,但对于生产环境来说太慢。最后两个阶段会转换模型,并将其部署在 API 之后。

导出为 ONNX / TensorRT

nemotron embed export -c default

这会将微调后的检查点导出为 ONNX(opset 17)。可选地,它会编译一个 TensorRT 引擎以实现最大推理吞吐量,并可为批量大小(1–64)和序列长度(3–256)配置优化配置文件:

# ONNX only (runs anywhere)
nemotron embed export -c default export_to_trt=false

# FP8 quantization for further speedup
nemotron embed export -c default quant_cfg=fp8

使用 NVIDIA NIM 部署

导出的模型部署在 NVIDIA NIM 容器内——这是一个生产就绪的推理微服务,暴露 OpenAI 兼容的 /v1/embeddings 端点:

nemotron embed deploy -c default

运行后,任何客户端都可以调用它:

curl -X POST http://localhost:8000/v1/embeddings \
  -H "Content-Type: application/json" \
  -d '{"input": ["What cooling is needed for 8 H100 GPUs in a 2U chassis?"],
       "model": "custom",
       "input_type": "query"}'

由于 NIM 提供 OpenAI 兼容的 API,你可以将其直接接入任何使用该 embeddings API 格式的现有 RAG 流程——无需更改代码。

验证部署准确率

该流程包含一个 NIM 准确率验证步骤,会对已部署的端点运行相同的 BEIR 评估:

nemotron embed eval -c default eval_nim=true eval_base=false

这能捕捉 ONNX/TensorRT 转换带来的任何准确率损失。在容差范围内匹配的指标(@1 为 0.03,@5+ 为 0.01)会标记为通过;超出转换噪声的偏差会被标记出来。

整合起来

完整的嵌入微调流水线可以通过六条命令运行,从原始文档到部署模型。

# 1. Generate synthetic training data from your documents
nemotron embed sdg -c default corpus_dir=./data/my_docs

# 2. Prepare the training data (split data, mine hard negatives, unroll)
nemotron embed prep -c default

# 3. Fine-tune the embedding model
nemotron embed finetune -c default

# 4. Evaluate the base vs. fine-tuned model
nemotron embed eval -c default

# 5. Export the optimized model
nemotron embed export -c default

# 6. Deploy the model
nemotron embed deploy -c default

预计时间和资源

阶段 需要GPU? 预计时间 说明
SDG 否(使用API) 约1小时 因语料库大小和API速率限制而异
数据准备 是(40 GB显存) 约5分钟 在GPU上进行难负例挖掘
微调 是(80 GB显存) 约1小时 因数据集大小和训练轮数而异
评估 是(40 GB显存) 约5分钟
导出 是(40 GB显存) 约5分钟 TensorRT需要NGC容器
部署 是(40 GB显存) 约5分钟 NIM容器启动

总计:不到一天,其中大部分时间是无须人工干预的训练。对于小型语料库(约500份文档),整个流水线大约在2–3小时内完成。

该流水线可以端到端运行,但每个阶段也可以根据你的起点独立执行。例如,如果你有原始文档,可以从合成数据生成(SDG)开始;而已经包含难负例的数据集可以跳过前面的步骤,直接进行微调。由于每个阶段都使用JSON、BEIR和ONNX等标准格式,因此很容易集成自定义组件,或在其他工作流中复用中间输出。该配方在运行方式上也很灵活,支持在本地机器、Docker容器内或基于Slurm的集群上执行。

自己动手试试

如果你有领域文档并且手头有些时间,今天就可以生成你的第一批合成训练数据!完整流水线——从文档到已部署的领域适配嵌入模型——在单块GPU上不到一天即可运行。你可以从我们现成的nvidia/Retrieval-Synthetic-NVDocs-v1数据集开始,立即试用该流水线。告诉我们你构建了什么。
如果你觉得有用,请为Nemotron、NeMo Data Designer和NeMo Automodel点亮仓库星标。

来源:Hugging Face Blog · hf-mirror.com