视觉多语言：mcdse-2b 模型介绍

今天，我将介绍一种新的实验性多语言嵌入模型，用于灵活的视觉文档检索。mcdse-2b-v1 (🤗) 基于 MrLight/dse-qwen2-2b-mrl-v1 训练，并使用 DSE 方法。

这个模型允许您嵌入页面/幻灯片截图，并使用自然语言进行查询。无论是表格、图表、示意图、图像还是文本，mcdse-2b-v1 都能将所有内容编码成一个单一的嵌入向量，从而无需传统的 OCR、文档布局分析、阅读顺序检测、分块、表格/公式提取等操作。

• 在 🇮🇹 意大利语、🇪🇸 西班牙语、🇬🇧 英语、🇫🇷 法语和 🇩🇪 德语上表现出色.

• 套娃表示学习：嵌入可以高效地从 1536 维缩放到 256 维。您可以将大小减少 6 倍，同时仍保持 95% 的嵌入质量。

• 二值化表现卓越：768 维二值向量保留了 1536 维浮点向量 99% 的检索质量。使用二值向量，您可以在 10GB 空间中编码 1 亿个多语言页面。

• 快速 vLLM 推理：在 vLLM 上运行推理，并高效地大规模提供嵌入服务，可用于生产环境。查看部署了解更多信息。

我的基准测试并非完美无缺，因此我鼓励您在自己的数据上测试模型。这是一个早期版本，还有很大的改进空间。尽管如此，结果仍突显了强大的多语言检索器，它能出色地适应各种内存/速度要求。

训练

mcdse-2b 基于 MrLight/dse-qwen2-2b-mrl-v1 使用低秩适配器 (LoRA) 在多语言文档语料库上进行训练。我使用 8xRTX3090 显卡，采用 DSE 方法，并使用以下参数进行了训练：

• 纪元 = 1
• 预热比例 = 0.1
• 学习率 = 1e-5
• 优化器 = adamw_torch
• 调度 = 线性
• 总批次大小 = 16
• LoRA
  • Alpha = 64
  • R = 16
  • Dropout = 0.1
  • DoRA = True

数据集

该数据集包含从公共互联网自动抓取的 24K PDF 页面。从每个文档中提取随机页面，转换为压缩 JPEG 图像，并过滤以去除空白页和重复项。最终的页面截图是独一无二的，涵盖了广泛的主题。

我使用 gemini-flash-1.5-002 根据每张图片生成查询。Gemini 被指示生成三种类型的查询：

• 一个宽泛的主题查询：总结文档的整体主题。
• 一个具体的详细问题：捕捉内容中的细微差别。
• 一个视觉查询：侧重于图表、图形、图像或签名等视觉元素。

整个训练和评估数据集仅花费 2 欧元生成（感谢 Gemini Flash！）

然后，每张图像根据其文本密度进行 0 到 2 的分类。我使用了在 DocLayNet 上微调的 omoured YOLOv10n 模型，用于检测图形与文本等区域。根据这些比例，我启发式地计算文本密度。我计划使用这种分类来提高模型在文本密集文档上的性能。

• 0 = 仅视觉元素
• 1 = 视觉元素和文本混合
• 2 = 仅文本

评估和训练数据集尚未发布。我非常愿意开源它们，但我仍然不确定如何在不违反任何许可（如果有的话）的情况下正确地做到这一点。如果您知道如何帮助我，请联系我！

训练运行

模型按以下顺序依次训练每种语言：

1. 法语：6k 样本
2. 西班牙语：6k 样本
3. 意大利语：6k 样本
4. 德语：6k 样本

这个顺序是根据基础模型在这些语言中的检索性能确定的，第一个是表现最好的。我的直觉是，鉴于数据集较小，从更强的语言开始可能有助于平衡模型在整体上的改进。

在达到这个最终检查点之前，我进行了多次运行，以测试各种策略并验证我的一些直觉。

• 语言顺序： 我交换了最后两种语言的顺序，发现最后训练德语可以使其在评估中的性能提高 1.7%，同时保持其他语言相似的分数。

• 模型初始化： 我使用每种语言的 10k mmarco 对初始化模型。这导致所有语言的性能更差，特别是低维嵌入。例如，使用 512 维嵌入的法语 NDCG@5 在使用 mmarco 训练时下降了 2%。

• 不同的图像调整算法： 我开发了一个自定义调整函数（custom_resize），它在将图像缩放到 min_pixels 和 max_pixels 范围内时严格保留图像的宽高比。所有评估都使用 qwen_vl_utils 中的标准调整函数。使用自定义调整函数训练的模型优于标准方法，NDCG@5 平均提高了 +1.7%（1536 维）。探索使用此 custom_resize 函数训练 ColQwen 模型会很有趣。

  调整函数	平均	英语	意大利语	西班牙语	法语	德语
  qwen2_vl_utils	80.8	80.2	80.5	79.6	81	82.6
  自定义调整	82.2	80.8	81.2	80.7	84.5	83.8
  	+1.7%	+0.7%	+0.9%	+1.4%	+4.0%	+1.4%

评估

由于公开可用的多语言文档图像检索数据集有限，该模型已使用自定义数据集进行评估。此评估数据集专门设计用于基准测试模型在各种语言中的性能。

此评估数据集是使用与训练数据集相同的方法和流程创建的。但是，文档主题通常不同，并且训练和评估数据集之间没有共享图像，以避免任何评估污染。NDCG 分数是通过对每种语言的 1K 文档索引运行 100 个唯一查询来计算的。

套娃表示学习

此模型使用套娃表示学习（MRL）进行训练，维度分别为：1536、1024、768、512、384、256。训练过程中使用的损失函数经过校准，以跟踪所有这些维度上的性能，从而使模型能够预先加载最重要的识别信息。这有效地允许您根据自己的规模和预算缩小嵌入维度。

每个维度的平均 NDCG@5。有趣的是，该模型在英语方面显示出改进，尽管该语言未包含在训练集中。该模型在 256 维度上性能提高了 6%，并且在所有维度上平均整体提高了 4%。评估使用 FAISS 的 IndexFlatL2 进行。

NDCG@5（浮点数）

二值嵌入

mcdse-2b-v1 在二值化方面表现明显更好，尤其是在较低维度上。该模型在 256 维度上提高了 23%，整体平均提高了 13%。评估使用 FAISS 的 IndexBinaryFlat 进行。但是，为什么二值嵌入更优呢？

此表显示，mcdse-2b-v1 的二进制嵌入比基础模型的 1536 维浮点向量性能高出 1%，同时内存消耗减少了 16 倍。除了这些优势，二进制嵌入还可以通过汉明距离进行 40 倍的快速搜索，因为比较两个二进制向量只需 2 个 CPU 周期（异或、位计数）。

NDCG@5（二进制）

ShiftProject

vidore/shiftproject_test 数据集是 ViDoRe 基准测试的一部分。它包含来自 Shift Project 关于环境的法语查询和文档。查询是使用 Claude-3 Sonnet 在与生成 vidore/colpali_train_set 抓取文档查询相同的法语翻译提示下生成的。

这是 ShiftProject 数据集上的 NDCG@5，使用 1536 浮点维度，最多评估 960 个图像补丁。

我原以为 mcdse-2b-v1 的分数会高于基础模型，结果却低了 3%。基础模型是在 colpali 训练集上训练的，我曾以为它可能对“Claude-3 Sonnet 风格”的查询过度优化了。为了调查这一点，我使用 gemini-1.5-flash-002 和我的提示重新生成了 ShiftProject 数据集查询。

分数变化很大，但在这种情况下，mcdse-2b-v1 提高了 5%。这些结果倾向于表明两种可能的原因：

1. 基础模型对“Claude-3 Sonnet 风格”的查询过度优化
2. 我的模型对“gemini-1.5-flash-002 风格”的查询过度优化

在这两种情况下，我相信 mcdse-2b-v1 通过理解更广泛的查询分布，减轻了这些过度优化。

我生成的 Gemini 查询有两种格式：问题和查询。colpali_train_set 生成的查询只有问题。我还测试了只使用 Gemini 查询和只使用 Gemini 问题的两种模型。

基础模型在 Gemini 问题上表现好 7%，在 Gemini 查询上表现差 7%。查询和问题之间的平均分数几乎相同（66.7 和 66.5）。这表明我的模型已经缓解了之前提到的过度优化，并且通常能够更好地理解更广泛的查询。在更多多语言数据上进行训练可能会提高这个平均值，并最终提高 ShiftProject 上的性能。

Cohere Embed v3 图像

我使用最近发布的 cohere embed-multilingual-v3.0 多模态 嵌入模型对较小版本的英语数据集进行了一些初步（且仓促）测试。该模型取得了 71 的 NDCG@5 分数，而 mcdse-2b-v1 的分数约为 84。我正在为该模型进行更全面的评估。

部署

在 HuggingFace Transformers 上，您可以使用 RTX3090 显卡以 32 的批处理大小（35TFLOPS）编码约 3 张图片/秒。像 RTX 4000 Ada 这样的更常见的推理端 GPU 大致能提供相同的吞吐量。

vLLM

vLLM 官方仅支持 Qwen2VL 用于生成，我添加了一个新的模型类 Qwen2VLForEmbeddingGeneration 来支持嵌入任务。在 vLLM 上运行推理应该比 HuggingFace Transformers 快约 5 倍。

下载新的模型类

git clone https://github.com/marplex/mcdse && cd mcdse

下载 mcdse-2b-v1 用于本地推理

from huggingface_hub import snapshot_download
snapshot_download(repo_id="marco/mcdse-2b-v1", local_dir="/path/to/model/mcdse-2b-v1")

编辑 config.json

将 Qwen2VLForConditionalGeneration 替换为 Qwen2VLForEmbeddingGeneration

sed -i -e 's/Qwen2VLForConditionalGeneration/Qwen2VLForEmbeddingGeneration/g' /path/to/model/mcdse-2b-v1/config.json

检查 vllm/main.py 进行本地推理

#vllm/main.py
from qwen2_vl_dse import Qwen2VLForEmbeddingGeneration, get_query_prompt, get_document_prompt
from vllm import ModelRegistry, LLM
from PIL import Image

ModelRegistry.register_model("Qwen2VLForEmbeddingGeneration", Qwen2VLForEmbeddingGeneration)

llm = LLM(
    model="/path/to/model/mcdse-2b-v1",
    limit_mm_per_prompt={
        "image": 1
    }
)

# Encode queries
query_prompt, image = get_query_prompt("Quali erano le passività totali al 31 dicembre 2017?")
outputs = llm.encode({"prompt": query_prompt, "multi_modal_data": {"image": [image]}})
outputs[0].outputs.embedding #1536 dimensional embedding

# Encode documents
dummy_document_image = Image.new('RGB', (256, 256))
document_prompt, image = get_document_prompt(dummy_document_image)
outputs = llm.encode({"prompt": document_prompt, "multi_modal_data": {"image": [image]}})
outputs[0].outputs.embedding #1536 dimensional embedding

结论

这是我第一次训练模型，虽然充满挑战，但也乐趣无穷。如果没有 HuggingFace 团队和贡献者的出色工作，我不可能做到这一点。我还要感谢 Manuel Faysse、Tony Wu 以及整个 Vidore 团队在 ColPali 上的工作，感谢 Xueguang Ma 在 Tevatron 代码库上的所有工作以及训练了一个非常强大的基础模型。我还受到了 Benjamin Clavié 及其令人印象深刻的模型公告的启发。

我希望这个模型对您的检索和 RAG 管道有用。正如开头所述，我的基准测试远非完美，实际场景中的结果可能会有所不同。我鼓励您在自己的用例中进行测试。总的来说，视觉检索的一个显著优势是，您可以通过简单地嵌入页面来抛弃复杂的索引管道。这就是未来！