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],同时使用其视觉编码器和文本编码器,并弃用文本解码器。
- 视觉编码器:利用 DaViT [4],从输入图像中生成 N 个视觉嵌入向量。
- 文本编码器:基于 BART [5] 编码器,它将 N 个视觉嵌入作为输入,生成上下文嵌入。然后,这些上下文嵌入通过一个线性层投影到 128 维向量,以优化嵌入的存储需求,类似于 ColBERT 和 ColPali。
检索系统
文档检索通常涉及两个主要步骤:
索引:机器学习模型为每个文档生成嵌入,捕获关键特征。这些嵌入存储在索引中,以便后续高效检索。
查询:用户的查询被编码成嵌入,然后将其与索引的文档嵌入进行比较,以根据相似性检索和排序结果。
传统的检索模型为每个文档/查询生成单个嵌入向量。相比之下,ColBERT风格的模型(如ColPali)使用上下文后期交互,将文档和查询编码成上下文嵌入袋。相似性通过MaxSim操作在这些嵌入袋之间计算。
ColFlor 遵循这种方法。在索引过程中,图像和(Florence-2 论文中使用的)一个
训练设置
我们从 Florence-2-base 初始化模型权重,新线性投影层除外,该层是随机初始化的。最初,训练不稳定,尽管进行了一些超参数搜索,损失仍无法收敛。为了解决这个问题,我们首先移除了随机初始化的投影层,并通过直接在文本编码器输出嵌入上应用 MaxSim 操作,对模型进行了 5 个 epoch 的训练。这稳定了训练并改善了收敛。之后,我们重新引入了线性层,并在 ViDoRe 数据集 [1] 上对模型进行了 40 个 epoch 的微调,使用 2e-5 的学习率和 64 的批大小,在 4 块 A100 GPU 上进行。
评估
我们对 ColFlor 在 ViDoRe 基准上进行了评估,该基准包含 10 个文档检索任务子类别。我们将其分组如下:
文本丰富的英文文档:包括 DocVQA、TatDQA 等学术数据集,以及人工智能、能源、政府报告和医疗保健等真实世界实用数据。
图表文档:包括 InfoVQA 和 ArxivQA,主要由图表、示意图和信息图表等复杂视觉内容组成。
法文文档:包括 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).