优化 ColPali 以实现大规模检索，速度提升 13 倍

ColPali 是文档检索领域的一个迷人飞跃。它在处理视觉丰富的 PDF 文件方面表现出惊人的精度，但将其扩展到处理真实世界数据集时，会带来一系列计算挑战。

我们解决了这些挑战，使 ColPali 的速度提升了 13 倍，同时没有牺牲其闻名的精度。

扩展困境

ColPali 仅针对一页 PDF 就会生成 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{ 次比较!} $$

这意味着需要数万亿次比较才能构建索引。即使是 HNSW 等先进的索引算法，也难以应对这种规模，因为计算成本与每页多向量的数量呈二次方增长。

我们转向了一种混合优化策略，结合了池化（以减少计算开销）和重排（以保持准确性）。

在我们深入探讨之前，请观看我们的网络研讨会视频，了解完整的演示演练。

对于那些渴望探索的人，代码库可在此处获取。

两阶段检索过程

池化

池化在机器学习中是一种众所周知的数据压缩方法，同时保留了重要信息。对于 ColPali，我们通过对文档 32x32 网格中的行进行池化，将每页 1,030 个向量减少到仅 38 个向量。

最大池化和平均池化是两种最流行的类型，因此我们决定在网格的行上测试这两种方法。同样，我们可以在列上应用池化，这将在未来进行探索。

• 平均池化： 对行中的值取平均。
• 最大池化： 为每个特征选择最大值。

32 个向量代表池化后的行，而 6 个向量编码从 ColPali 特殊标记（例如，用于序列的开头，以及“描述图像”等任务特定指令）中派生出的上下文信息。

在我们的实验中，我们选择保留这 6 个额外的向量。

“将 ColPali 作为重排器”实验

池化极大地降低了检索成本，但存在损失细粒度精度的风险。为了解决这个问题，我们实施了一个两阶段检索系统，其中使用 ColPali 生成的嵌入通过网格行进行最大/平均池化，以创建用于初始检索阶段的轻量级向量，然后使用原始高分辨率嵌入进行重排

1. 池化检索： 使用轻量级池化嵌入快速检索前 200 个候选。
2. 完全重排： 使用原始高分辨率嵌入对这些候选进行精炼，最终得到前 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 倍。

度量

平均池化与原始 ColPali 保持了几乎相同的质量，NDCG@20 = 0.952，Recall@20 = 0.917。最大池化表现不佳，不足以被认为是可行的，因为它牺牲了显著的准确性，却没有带来显著的速度优势。

下一步是什么？

未来的实验可以进一步推动这些结果

• 研究列式池化以实现额外压缩。
• 测试半精度 (float16) 向量以平衡内存使用和速度。
• 在预取期间跳过特殊多向量以简化检索。
• 将量化与过采样相结合，以实现更快的搜索。

亲自尝试

想亲眼看看它的效果吗？探索完整的代码库并尝试 ColPali 优化

• 演示 Notebook： GitHub 存储库
• 网络研讨会演练： 在此观看

加入社区并分享您的结果！