ColFlor：迈向BERT大小的视觉-语言文档检索模型

引言

2024年6月，ColPali [1] 作为一款无需OCR的文档检索模型被引入，其基于PaliGemma [2] 构建，通过直接处理图像而非使用易出错且资源消耗大的OCR管道，改变了PDF文档检索的范式。然而，ColPali拥有30亿参数，可能在计算上成本高昂，特别是对于大型文档数据库。相比之下，ColBERT等文本检索模型效率更高，仅有数亿参数，但它们需要易出错且昂贵的OCR管道。为了弥补这一差距，我们引入了ColFlor，这是一款无需OCR的视觉文档检索模型，参数仅为1.74亿。ColFlor比ColPali小17倍，查询编码速度快9.8倍，图像编码速度快5.25倍，而在文本丰富的英文文档上性能仅下降1.8%。

建模

架构

ColFlor 利用 Florence-2 的架构 [3]，同时使用其视觉编码器和文本编码器，并弃用文本解码器。

1. 视觉编码器：利用 DaViT [4]，从输入图像中生成 N 个视觉嵌入向量。
2. 文本编码器：基于 BART [5] 编码器，它将 N 个视觉嵌入作为输入，生成上下文嵌入。然后，这些上下文嵌入通过一个线性层投影到 128 维向量，以优化嵌入的存储需求，类似于 ColBERT 和 ColPali。

检索系统

文档检索通常涉及两个主要步骤：

1. 索引：机器学习模型为每个文档生成嵌入，捕获关键特征。这些嵌入存储在索引中，以便后续高效检索。

2. 查询：用户的查询被编码成嵌入，然后将其与索引的文档嵌入进行比较，以根据相似性检索和排序结果。

传统的检索模型为每个文档/查询生成单个嵌入向量。相比之下，ColBERT风格的模型（如ColPali）使用上下文后期交互，将文档和查询编码成上下文嵌入袋。相似性通过MaxSim操作在这些嵌入袋之间计算。

ColFlor 遵循这种方法。在索引过程中，图像和（Florence-2 论文中使用的）一个 文本标记都被编码，以为输入图像创建文本感知上下文嵌入。在查询时，文本编码器处理查询以生成文本嵌入。最后，使用 MaxSim 操作测量图像和查询嵌入之间的相似度。

训练设置

我们从 Florence-2-base 初始化模型权重，新线性投影层除外，该层是随机初始化的。最初，训练不稳定，尽管进行了一些超参数搜索，损失仍无法收敛。为了解决这个问题，我们首先移除了随机初始化的投影层，并通过直接在文本编码器输出嵌入上应用 MaxSim 操作，对模型进行了 5 个 epoch 的训练。这稳定了训练并改善了收敛。之后，我们重新引入了线性层，并在 ViDoRe 数据集 [1] 上对模型进行了 40 个 epoch 的微调，使用 2e-5 的学习率和 64 的批大小，在 4 块 A100 GPU 上进行。

评估

我们对 ColFlor 在 ViDoRe 基准上进行了评估，该基准包含 10 个文档检索任务子类别。我们将其分组如下：

1. 文本丰富的英文文档：包括 DocVQA、TatDQA 等学术数据集，以及人工智能、能源、政府报告和医疗保健等真实世界实用数据。

2. 图表文档：包括 InfoVQA 和 ArxivQA，主要由图表、示意图和信息图表等复杂视觉内容组成。

3. 法文文档：包括 TabFQuAD 和 Shift，用于测试模型的多语言能力。

结果（如下表所示）表明，ColFlor 在文本丰富的英文文档上与 ColPali 表现相当，尽管其尺寸显著更小，平均性能仅下降 1.8%。值得注意的是，ColFlor 在 TatDQA（一个源自公开真实世界财务报告的 VQA 数据集）和 Health 数据集上的表现优于 ColPali。这突显了 ColFlor 在实际应用中的潜力及其高效扩展的能力。在图表文档类别中，性能差距更为明显，这可能是由于基础模型（Florence-2）侧重于文本丰富的文档，并且在图表上的训练有限。我们计划在未来通过继续对 Florence-2 进行图表预训练来解决这个问题。最后，ColFlor 在法文文档上的表现不佳，因为 Florence-2 仅为英文设计，缺乏多语言支持。

ColFlor的效率

ColFlor 模型旨在为 ColPali 提供一个高效、经济且性能卓越的替代方案，使新的无 OCR 文档检索范式能够服务于计算资源有限（GPU 贫乏）的用户。ColFlor 在参数方面比 ColPali 小 17 倍。我们使用 float32 数据类型在一台免费 T4 GPU 上对两个模型的前向传播进行了基准测试。对于图像编码，ColFlor 使用 32 的批次大小，ColPali 使用 2 的批次大小。对于查询编码，我们使用 1 的批次大小来模拟在线查询。如下图所示，ColFlor 的图像编码速度快 5.25 倍，查询编码速度快 9.8 倍。此外，ColFlor 以更高的分辨率（768x768 对比 ColPali 的 448x448）处理图像，同时生成更少的上下文嵌入（587 对比 1024），从而降低了存储成本。

结论与未来工作

我们介绍了 ColFlor，一个用于无 OCR 文档检索的 BERT 大小模型。ColFlor 比 ColPali 小 17 倍，图像编码速度快 5.25 倍，查询编码速度快 9.8 倍，而在文本丰富的英文文档上的性能仅下降 1.8%。展望未来，我们正在探索 Florence-2 的持续预训练，以增强图表理解能力。我们还计划开发 Florence-2 的多语言变体，使 ColFlor 能够支持更广泛的语言，以实现更广泛的应用。

在这篇博客文章中，我们分享了我们关于多模态文档检索研究的初步发现和早期成果。我们对这项工作在开源社区中的潜力感到兴奋。此外，我们目前正在努力进行进一步的改进，并撰写一份更全面的技术报告！。

资源

🧠 模型: https://huggingface.co/ahmed-masry/ColFlor
💻 演示: https://huggingface.co/spaces/ahmed-masry/ColFlor-Demo
🏋️‍♂️ 训练代码: https://github.com/AhmedMasryKU/colflor
📊 评估代码: https://github.com/AhmedMasryKU/vidore-benchmark-colflor

如果您对此工作有任何疑问，请随时通过 masry20@yorku.ca 联系 Ahmed Masry

致谢

这项工作在加拿大约克大学的智能可视化实验室进行。得到了加拿大自然科学与工程研究委员会 (NSERC) 和加拿大创新基金会 (CFI) 的支持。此外，它还获得了谷歌 PaliGemma 学术计划提供的 GCP 积分奖励。

我们非常感谢 ColPali 团队提供的详细记录的训练和评估 GitHub 仓库，它们对我们的模型开发至关重要。

引用

如果您计划在您的研究中使用 ColFlor，请考虑以下方式引用我们：

@misc{masry2024colflor,
    title={ColFlor: Towards BERT-Size Vision-Language Document Retrieval Models},
    url={https://huggingface.co/blog/ahmed-masry/colflor},
    author={Masry, Ahmed},
    month={October},
    year={2024}
}

参考文献

[1] Faysse, M., Sibille, H., Wu, T., Viaud, G., Hudelot, C. and Colombo, P., 2024. ColPali: Efficient Document Retrieval with Vision Language Models. arXiv preprint arXiv:2407.01449.

[2] Beyer, L., Steiner, A., Pinto, A.S., Kolesnikov, A., Wang, X., Salz, D., Neumann, M., Alabdulmohsin, I., Tschannen, M., Bugliarello, E. and Unterthiner, T., 2024. PaliGemma: A versatile 3B VLM for transfer. arXiv preprint arXiv:2407.07726.

[3] Xiao, B., Wu, H., Xu, W., Dai, X., Hu, H., Lu, Y., Zeng, M., Liu, C. and Yuan, L., 2024. Florence-2: Advancing a unified representation for a variety of vision tasks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 4818-4829).

[4] Ding, M., Xiao, B., Codella, N., Luo, P., Wang, J. and Yuan, L., 2022, October. Davit: Dual attention vision transformers. In European conference on computer vision (pp. 74-92). Cham: Springer Nature Switzerland.

[5] Lewis, M., 2019. Bart: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension. arXiv preprint arXiv:1910.13461.

[6] Khattab, O. and Zaharia, M., 2020, July. Colbert: Efficient and effective passage search via contextualized late interaction over bert. In Proceedings of the 43rd International ACM SIGIR conference on research and development in Information Retrieval (pp. 39-48).