Llama-NeMoRetriever-ColEmbed:英伟达最先进图文检索技术的开发者指南
社区文章 2025年7月9日发布
对跨文本和图像模态的强大检索系统的需求正在迅速增长。Llama-NemoRetriever-ColEmbed 系列引入了一种统一的文本-图像检索方法,在多项基准测试中取得了最先进的结果。本文将从技术角度概述模型架构、训练策略、评估结果和实际权衡,重点关注对开发者最重要的方面。
模型架构
带有后期交互的双编码器
- 基础: 模型架构基于 NVIDIA 的 Eagle2 视觉语言模型 (VLM),将因果注意力替换为双向注意力。
- 动态图像平铺: 支持可变输入分辨率,通过
max_input_tiles
和min_input_tiles
参数控制。 - ColBERT 风格的后期交互: 不再将序列压缩成单个向量,而是每个查询令牌嵌入与文档的所有令牌嵌入通过 MaxSim 运算符进行交互。这实现了细粒度的令牌级匹配。
模型变体 | 参数 (B) | 嵌入维度 |
---|---|---|
1B | 2.42 | 2048 |
3B | 4.41 | 3072 |
训练流程
两阶段训练
阶段1:仅文本预训练
- 模型使用对比损失在大型纯文本检索数据集上进行训练。
- 为文本建立强大的语义表示。
阶段2:文本-图像微调
- 在多样化的文本-图像对上进行微调,将文本和视觉表示对齐到共享的嵌入空间中。
数据集
- 纯文本: HotpotQA、MIRACL、自然问题、Stack Exchange、SQuAD、Tiger Math/Stack。
- 文本-图像: ColPali、Wiki-SS-NQ、VDR、VisRAG-Ret-Train-Synthetic、VisRAG-Ret-Train-In-domain、Docmatix。
评估结果
基准测试
- ViDoRe V1 & V2: 3B 模型在 nDCG@5 分数上分别达到 91.0 (V1) 和 63.5 (V2),截至 2025 年 6 月,在两个排行榜上均处于领先地位。
- MTEB 视觉文档检索: 3B 模型得分为 83.1,优于更大的 7B 模型。
- MIRACL-VISION: 展现了强大的多语言检索能力,其中 3B 变体取得了最高的综合平均分数 (0.5841)。
模型 | 参数 | 嵌入维度 | MTEB VDR | ViDoRe V1 | ViDoRe V2 |
---|---|---|---|---|---|
nvidia/llama-nemoretriever-colembed-1b-v1 | 2B | 2048 | 82.63 | 90.5 | 62.1 |
nvidia/llama-nemoretriever-colembed-3b-v1 | 4B | 3072 | 83.10 | 91.0 | 63.5 |
系统权衡
存储和延迟
- 后期交互模型: 需要存储所有令牌嵌入,导致显著的存储开销。例如,一个 3B 模型,如果使用 3072 维嵌入,一百万张图片需要超过 10 TB 的存储空间。
- 双编码器模型: 每个文档只存储一个向量,对于相同规模的语料库,仅需几 GB 存储空间。
- 降维: 线性投影层可以将存储空间减少高达 88%,且准确率损失极小。
检索流程选择
- 后期交互: 准确率更高,但存储和延迟也更高。
- 双编码器 + 重排序器: 存储空间更小,通过重排序可以达到有竞争力的准确率,但每次查询的推理时间会增加。
架构 | 存储 (100万张图片, GB) | ViDoRe V1 | ViDoRe V2 | 额外延迟 (毫秒/查询) |
---|---|---|---|---|
ColEmbed 3B (3072d) | 10,311.1 | 0.9106 | 0.6357 | 不适用 |
ColEmbed 3B (512d) | 1,230.2 | 0.9064 | 0.6109 | 不适用 |
双编码器 llama-vlm-embed-v1 (2048d)*¹ | 3.8 | 0.8313 | 0.5178 | 不适用 |
双编码器 llama-vlm-embed-v1 + 重排序器**¹ | 3.8 | 0.9064 | 0.6214 | 2,368 |
- *一种商业多模态检索模型,将用户查询表示为文本,文档表示为图像。
- **结果是使用内部开发的 VLM 重排序器对前 25 个文档进行排序后获得的。
- ¹ 这些数字可能略有不同,因为我们在使用 mteb 包计算排行榜结果之前,使用不同的代码库在 Vidore V1 和 V2 数据集上进行了评估。有关更多详细信息,请参阅完整的技术报告。
实际考量
- 部署: 根据您的存储、延迟和准确性要求选择模型和架构。
- 小语料库、高查询量: 可能更倾向于不带重排序器的较大嵌入模型。
- 大语料库、中等查询量: 带重排序器的小型嵌入模型可能更具成本效益。
- 向量数据库支持: 后期交互模型需要对令牌级相似度搜索的专门支持。
Llama-NemoRetriever-ColEmbed 代表了可扩展、高性能文本-图像检索的重大进步,在 ViDoRe V1、ViDoRe V2 和 MIRACL-VISION 基准测试中取得了最先进的结果。结合大规模文本和图像数据的两阶段训练流程,实现了强大的泛化能力和多语言检索能力。1B 和 3B 模型变体的发布为未来的多模态检索场景的研究和实际部署奠定了坚实基础。
如需更深入的技术理解和全面的权衡分析,请阅读完整研究报告。如果您有兴趣试验这些模型,请直接访问build.nvidia.com/explore/retrieval 试用 NeMo Retriever 模型。这是在您自己的应用程序和工作流程中探索最先进检索的绝佳机会。