优化 ColPali 实现大规模检索，结果速度提升 13 倍
Evgeniya Sukhodolskaya, Sabrina Aquino
·2024 年 11 月 27 日
ColPali 是文档检索领域一个引人入胜的飞跃。它在处理视觉丰富 PDF 方面的精确度非常惊人，但将其扩展到处理真实世界数据集会带来一系列计算挑战。
以下是我们如何解决这些挑战，使 ColPali 速度提升 13 倍，同时不牺牲其闻名的精确度。
规模化困境
对于 一页 PDF，ColPali 会生成 1,030 个向量。 虽然这对于小规模任务来说易于管理，但在需要存储数十万个 PDF 的真实世界生产环境中，规模化挑战变得尤为重要。
考虑以下场景
- 数据集大小： 20,000 页 PDF。
- 向量数量： 每页生成约 1,000 个 128 维向量。
比较总数计算如下：
$$ 1,000 \cdot 1,000 \cdot 20,000 \cdot 128 = 2.56 \times 10^{12} \text{ comparisons!} $$
建立索引需要数万亿次比较。即使是像 HNSW 这样的高级索引算法也难以应对这种规模，因为计算成本随每页多向量的数量呈二次方增长。
我们转向了一种混合优化策略，结合了池化（用于减少计算开销）和重排（用于保持准确性）。
在我们深入探讨之前，请观看我们的网络研讨会视频，了解完整的演示过程。
对于那些渴望探索的人，此处提供代码库。
两阶段检索过程
池化
池化在机器学习中是一种众所周知的数据压缩方法，同时保留重要信息。对于 ColPali，我们通过对文档 32x32 网格中的行进行池化，将每页的 1,030 个向量减少到仅 38 个向量。
最大池化和平均池化是两种最流行的类型，因此我们决定在网格的行上测试这两种方法。同样，我们也可以对列应用池化，这我们计划在未来探索。
- 平均池化： 计算行上的平均值。
- 最大池化： 为每个特征选择最大值。
32 个向量代表池化后的行，而 6 个向量编码了 ColPali 特殊 token 中派生的上下文信息（例如，
在我们的实验中，我们选择保留这 6 个附加向量。
“ColPali 作为重排器”实验
池化极大地降低了检索成本，但存在丢失细粒度精确度的风险。为了解决这个问题，我们实现了一个两阶段检索系统，其中使用 ColPali 生成的嵌入按网格行进行最大/平均池化，以创建用于初始检索阶段的轻量级向量，然后使用原始高分辨率嵌入进行重排。
- 池化检索： 使用轻量级池化嵌入快速检索前 200 个候选结果。
- 完整重排： 使用原始高分辨率嵌入对这些候选结果进行精炼，生成最终的前 20 个结果。
实现
我们通过合并以下内容创建了一个包含 20,000 多页独特 PDF 的自定义数据集：
- ViDoRe 基准： 专为 PDF 文档检索评估而设计。
- UFO 数据集： 视觉丰富的文档与 由 Daniel van Strien 生成的合成查询配对。
- DocVQA 数据集： 大量源自文档的问答对。
每个文档被处理成 32x32 的网格，生成全分辨率和池化嵌入。全分辨率嵌入包含每页 1,030 个向量，而池化嵌入则包括平均池化和最大池化变体。
所有嵌入都存储在 RAM 中，以避免在检索速度实验期间出现缓存效应。
实验设置
我们使用 1,000 个查询评估了检索质量。首先，池化嵌入检索了前 200 个候选结果。然后，全分辨率嵌入对它们进行了重排，生成了最终的前 20 个结果。
为了衡量性能，我们使用了
- NDCG@20： 衡量排名质量（前 20 个结果与预期的一致程度）。
- Recall@20： 衡量此方法与原始 ColPali 检索之间的重叠度。
结果
实验显示在速度和准确性方面有显著改进。与单独使用全分辨率嵌入相比，检索时间提升了 13 倍。
指标
| 池化类型 | NDCG@20 | Recall@20 |
|---|---|---|
| 平均 | 0.952 | 0.917 |
| 最大 | 0.759 | 0.656 |
平均池化保留了与原始 ColPali 几乎相同的质量，NDCG@20 = 0.952，Recall@20 = 0.917。最大池化表现不佳，无法被视为可行，因为它牺牲了显著的准确性而没有带来有意义的速度优势。
下一步是什么？
未来的实验可以进一步提升这些结果
- 研究按列池化以实现额外压缩。
- 测试半精度 (float16) 向量以平衡内存使用和速度。
- 在预取期间跳过特殊多向量以简化检索。
- 将量化与过采样相结合以实现更快的搜索。
亲自尝试
好奇想看看实际效果？探索完整的代码库并进行 ColPali 优化实验
加入社区并分享您的结果！
