视觉文档检索走向多语言

Marco Cimolai

marco

关注

llamaindex

Logan Markewich

cheesyFishes

关注

llamaindex

摘要: 我们推出了 vdr-2b-multi-v1，这是目前最优秀的用于视觉文档检索的多语言嵌入模型。我们还发布了其仅支持英文的孪生模型 vdr-2b-v1，并开源了新的 vdr-multilingual-train 数据集。该数据集包含 50 万个高质量样本，是目前最大的用于视觉文档检索的开源多语言合成数据集。

隆重介绍 vdr-2b-multi-v1 (🤗)，一款专为跨多种语言和领域的视觉文档检索而设计的多语言嵌入模型。该模型旨在将文档页面的截图编码为密集的单向量表示，这将有效地支持对视觉丰富的多语言文档进行搜索和查询，而无需任何 OCR、数据提取流程、分块等操作……

vdr-2b-multi-v1 模型基于 MrLight/dse-qwen2-2b-mrl-v1，并使用我们自制的大量多语言查询-图像对数据集进行训练。该模型与 LlamaIndex 合作构建，是 mcdse-2b-v1 的下一代版本。我们的 vdr-2b-multi-v1 扩展并改进了训练时所用的学习和方法，从而产生了一个更强大、更优秀的模型。

• 训练数据覆盖 🇮🇹 意大利语、🇪🇸 西班牙语、🇬🇧 英语、🇫🇷 法语和 🇩🇪 德语： 这些语言共同构成了一个新的、大型的、开源的多语言训练数据集，包含 50 万个高质量样本。

• 低显存占用与更快的推理速度：在合成视觉文档检索（ViDoRe）基准测试中，我们仅支持英语的模型使用 768 个图像补丁的表现优于使用 2560 个图像补丁的基础模型。这使得推理速度提升了 3 倍，并显著降低了显存使用量。

• 跨语言检索：在真实世界场景中表现显著更优。例如，您可以使用意大利语查询来搜索德语文档。

• 套娃表示学习（Matryoshka Representation Learning）：您可以将向量大小缩减 3 倍，同时仍保持 98% 的嵌入质量。这在降低存储成本的同时，显著提高了检索速度。

使用方法

🎲 立即试用 vdr-2b-multi-v1，请访问此 Hugging Face Space！

借助 SentenceTransformers 和 LlamaIndex 的直接集成，使用 vdr-2b-multi-v1 生成嵌入比以往任何时候都更容易。只需几行代码即可开始

pip install -U llama-index-embeddings-huggingface

from llama_index.embeddings.huggingface import HuggingFaceEmbedding

model = HuggingFaceEmbedding(
    model_name="llamaindex/vdr-2b-multi-v1",
    device="cpu",  # "mps" for mac, "cuda" for nvidia GPUs
    trust_remote_code=True,
)

image_embedding = model.get_image_embedding("image.png")
query_embedding = model.get_query_embedding("Chi ha inventato Bitcoin?")

from sentence_transformers import SentenceTransformer

model = SentenceTransformer(
    model_name_or_path="llamaindex/vdr-2b-multi-v1",
    device="cuda",
    trust_remote_code=True,
    # These are the recommended kwargs for the model, but change them as needed if you don't have CUDA
    model_kwargs={
        "torch_dtype": torch.bfloat16, 
        "device_map": "cuda:0", 
        "attn_implementation": "flash_attention_2"
    },
)

embeddings = model.encode("image.png")

训练数据集

为视觉文档检索训练优秀的单向量模型需要高质量的数据，但目前现有的多模态数据集非常稀少且非多语言。

因此，我们花费了大量时间从头开始构建。原始数据集包含 50 万个多语言查询-图像样本，这些样本是使用公开的互联网 PDF 从头收集和生成的。与每个图像关联的查询是合成的，并使用 VLM 生成。相比之下，我们的数据集样本量比之前最大的用于多模态视觉文档检索的开源合成数据集（即为 ColPali 训练数据集 生成的爬取文档）多 10 倍。

数据收集

对于每种语言，我们都会生成一个涵盖许多不同主题的长搜索查询列表，然后用这些查询来搜索 PDF。我们使用搜索引擎的语言过滤功能来爬取仅包含指定语言的文档。这种“按主题搜索”的技术确保了模型能够接触到大量多样化的主题和领域，并且在现实生活场景中表现良好。

爬取过程产生了约 5 万份多语言文档。与之前 mcdse-2b-v1 模型中使用的方法不同，页面不是随机提取的。相反，每个 PDF 的每一页都通过一个文档布局分析模型来确定该页面是包含更多文本元素还是视觉元素。结果是一个数字，将页面分类为纯文本、纯视觉或混合型。然后，这个标注步骤被用来抽样约 10 万个页面，确保它们按页面类型均匀分布。

合成数据生成

然后使用 gemini-1.5-pro 和 Qwen2-VL-72B 生成查询。它们的任务是提出一个具体问题和一个泛化问题。之后只有具体问题用于训练模型，但强制 LLM 区分这两种问题通常会产生更强的具体问题，以用于信息检索训练。

生成后，进一步的清理步骤确保问题足够好用于训练。这包括：

• 确保语言正确
• 修复格式问题
• 移除 markdown
• 确保只提出一个问题
• 移除定位短语（例如“根据图1”、“本文档”……）

筛选与困难负样本挖掘

这个清理步骤确保了查询在语法上是正确的，并且遵循一些严格的准则。但这仍然不能保证这些查询足够好用于信息检索。

为了过滤掉不好的问题，我们使用 voyage-3 嵌入模型对每个泛化查询进行嵌入和索引。对于每个具体问题，我们在索引中搜索。如果其相关的泛化问题出现在前 100 个结果中，则该查询被标记为“好”。这种方法可以去除低熵、重复或过于相似的问题。平均而言，每种语言的数据集中有 40% 的查询被移除。

困难负样本随后仅使用 voyage-3 在具体问题上进行挖掘，固定阈值为 0.75。我们还进行了实验，使用了 nvidia/NV-Retriever-v1 中描述的正样本感知负样本挖掘方法，但在该数据集上，这种方法似乎产生了过于简单/遥远的负样本。

下载

(vdr-multilingual-train 🤗) 训练数据集现在已开源，并可直接在 Hugging Face 上获取。训练数据集包含 496,167 个 PDF 页面，其中只有 280,679 个与过滤后的查询相关联（使用上述方法）。那些没有查询的图像仍然被用作困难负样本。

该数据集由 5 个不同的子集组成，每种语言一个。您可以在这里直接浏览

或者，您可以通过在 load_dataset 中指定语言子集来单独下载各语言数据

from datasets import load_dataset

italian_dataset = load_dataset("llamaindex/vdr-multilingual-train", "it", split="train")

english_dataset = load_dataset("llamaindex/vdr-multilingual-train", "en", split="train")

french_dataset = load_dataset("llamaindex/vdr-multilingual-train", "fr", split="train")

german_dataset = load_dataset("llamaindex/vdr-multilingual-train", "de", split="train")

spanish_dataset = load_dataset("llamaindex/vdr-multilingual-train", "es", split="train")

评估

该模型已在 ViDoRe 基准测试和自定义构建的评估集上进行了评估，这些评估集能够测试其在纯文本、纯视觉和混合页面截图上的多语言能力。评估数据集也已在 Hugging Face 上公开（vdr-multilingual-test 🤗）。

我们确保了这些数据集中的任何页面都没有出现在训练集中，以避免任何评估污染。这些数据集是使用与训练数据集相同的方法收集和生成的，但样本量较小。过滤步骤全部是手动完成的：每个查询都经过评估、策划和改进（如果需要），以确保高质量的数据。

所有评估都是通过计算 NDCG@5 分数进行的，使用了 1536 维度的向量和可以用最多 768 个词元表示的图像分辨率。

多语言模型在每种语言和每种页面类型上都优于基础模型，平均提升了 2.3%。在 ViDoRe 基准测试中，它也表现略好（+0.5%）。我们微调的 vdr-2b-multi-v1 在性能上取得了巨大飞跃，特别是在非英语的纯视觉或混合页面上。例如，在德语纯视觉检索方面，NDCG@5 分数比基础模型提高了 6.33%。

我们还仅在英语子集上训练了一个版本（vdr-2b-v1 🤗）。在完整的 ViDoRe 基准测试中（使用 768 个图像词元进行评估），多语言版本和仅英语版本都优于基础模型。

更快的推理

仅支持英语的 vdr-2b-v1 模型在 ViDoRe 基准测试的合成数据集上也达到了与基础模型相当的性能，而只使用了 30% 的图像词元（768 vs. 2560）。这有效地带来了 3 倍的推理速度提升和更低的显存使用量。

跨语言检索

尽管该模型是针对每种语言单独训练的，但它在跨语言检索方面也有所改进。为了测试这种能力，我们将德语评估集的查询使用 DeepL 翻译成意大利语。文档页面的截图仍保持原始的德语。

该模型在所有文档类型上都显著更优，平均提升了 2.3%。这些检索能力对于现实世界的用例至关重要，特别是在像欧洲这样语言碎片化的大陆。例如，它能够在复杂的多语言源上进行与语言无关的搜索，如欧洲具有约束力的决定、说明手册、金融资产的 KID、药品说明书等等……

MRL 和二元嵌入

该模型使用套娃表示学习（Matryoshka Representation Learning，MRL）进行训练。训练期间使用的损失函数经过校准，以跟踪所有这些维度的性能，引导模型将最重要的识别信息前置。这实际上允许您根据您的规模和预算来缩减嵌入维度。要了解更多关于 MRL 的信息，Hugging Face 的这篇博文解释得非常好。

为了测试模型在不同向量维度下的检索能力，我们在意大利语->德语的跨语言基准测试中进行了评估。

NDCG@5 (浮点型)

NDCG@5 (二元)

1024 维度的浮点向量在质量和大小之间提供了非常好的平衡。它们的体积小了约 30%，但仍然保留了 99% 的检索性能。这对于 1536 维度的二元向量也是如此，它们每个向量的字节数减少了 10 倍，但仍保留了 97% 的检索质量。同样有趣的是，1536 维度的二元向量几乎与基础模型的 1536 维度浮点向量的性能相匹配。

结论与未来工作

我们相信 vdr-2b-multi-v1 和 vdr-2b-v1 将对许多用户证明其价值。

我们的多语言模型是同类产品中的首创，它显著提高了在多语言和跨语言场景中的性能，并且得益于 MRL 和二元量化，检索比以往任何时候都更高效、更快速。我们相信这将开启新的用例和机会，特别是在像欧洲这样语言碎片化的大陆。

其仅支持英语的孪生模型代表了对基础模型的重大改进，现在能够以 3 倍的速度嵌入文档，显存占用更少，并且具有相同（或更好）的检索质量。

所有这一切都得益于新的 vdr-multilingual-train 数据集。凭借 50 万个高质量样本，它是用于视觉文档检索的最大的多语言开源合成数据集。

未来的工作将探索我们的模型在适应新的和特定领域时的表现。这仍处于开发的早期阶段，在结果发表前还需要做更多的工作，但初步测试似乎已经表明，用极少的数据和计算资源就能实现令人印象深刻的检索增益。

敬请期待未来的更新！

链接

• 🎲 模型演示: Hugging Face Space
• 🤗 多语言模型: vdr-2b-multi-v1
• 🤗 仅英文模型: vdr-2b-v1
• 📂 训练数据集: vdr-multilingual-train
• 📂 评估数据集: vdr-multilingual-test