[论文解读] PLAID: An Efficient Engine for Late Interaction Retrieval

Keshav Santhanam, Omar Khattab+2 more|arXiv (Cornell University)|May 19, 2022

Topic Modeling被引用 8

一句话总结

PLAID 是一种高度优化的检索引擎，通过引入质心交互和质心剪枝，利用 ColBERTv2 加速晚期交互检索，在 GPU 上将搜索延迟降低高达 7 倍，在 CPU 上降低 45 倍，同时在最大规模达 14000 万段落的数据集上保持最先进水平的检索质量。

ABSTRACT

Pre-trained language models are increasingly important components across multiple information retrieval (IR) paradigms. Late interaction, introduced with the ColBERT model and recently refined in ColBERTv2, is a popular paradigm that holds state-of-the-art status across many benchmarks. To dramatically speed up the search latency of late interaction, we introduce the Performance-optimized Late Interaction Driver (PLAID). Without impacting quality, PLAID swiftly eliminates low-scoring passages using a novel centroid interaction mechanism that treats every passage as a lightweight bag of centroids. PLAID uses centroid interaction as well as centroid pruning, a mechanism for sparsifying the bag of centroids, within a highly-optimized engine to reduce late interaction search latency by up to 7$ imes$ on a GPU and 45$ imes$ on a CPU against vanilla ColBERTv2, while continuing to deliver state-of-the-art retrieval quality. This allows the PLAID engine with ColBERTv2 to achieve latency of tens of milliseconds on a GPU and tens or just few hundreds of milliseconds on a CPU at large scale, even at the largest scales we evaluate with 140M passages.

研究动机与目标

• 解决尽管晚期交互检索在大规模信息检索系统中表现最先进，但其高延迟的问题。
• 优化 ColBERTv2 的端到端检索流水线，该模型使用残差表示和基于质心的压缩技术。
• 在大规模场景下（特别是 1 亿段落以上）实现在 CPU 和 GPU 上的低延迟推理，尤其适用于生产部署。
• 在候选段落评分过程中显著降低计算和内存开销，同时保持检索质量。
• 开发一个模块化、高度优化的引擎，利用质心 ID 的离散特性实现高效的过滤和评分。

提出的方法

• 引入质心交互：将每段文本视为一组轻量级整数质心 ID，实现无需加载完整残差向量的快速近似评分。
• 使用预计算的质心-查询距离，避免在搜索过程中重复计算所有质心的相似度分数，从而实现在多个段落间复用。
• 实现质心剪枝：通过利用低分质心的稀疏性，在早期过滤阶段跳过距离较远的质心 ID，从而减少候选集大小。
• 为 GPU 设计无填充的 MaxSim 内核，为 CPU 实现优化的 C++ 内核，以加速残差解压缩和评分操作。
• 构建多阶段检索流水线，在完整残差评分前应用质心交互和剪枝，最小化对低质量候选段落的昂贵操作。
• 利用质心的固定离散词汇表，实现在 CPU 和 GPU 执行中高效的向量化操作和内存访问模式。

实验结果

研究问题

• RQ1 仅使用基于质心的表示是否足以在晚期交互检索中实现高召回率，从而在完整残差评分前实现高效过滤？
• RQ2 质心剪枝在不降低检索质量的前提下，能否有效减少候选段落数量？
• RQ3 算法优化（质心交互与剪枝）与底层内核优化相结合，能在多大程度上提升 CPU 和 GPU 上的端到端延迟？
• RQ4 PLAID 在不同硬件配置下，随着语料规模和检索深度（k）的增加，其扩展性如何？
• RQ5 在早期阶段用仅质心评分替代完整残差评分时，延迟降低与检索质量之间的权衡如何？

主要发现

• PLAID ColBERTv2 相较于原始 ColBERTv2，在 GPU 上实现最高 7 倍加速，在 CPU 上实现最高 45 倍加速，且在 MS MARCO v1、v2、Wikipedia 和 LoTTE 基准测试中未出现可测量的质量损失。
• 仅使用质心交互即可在 GPU 上实现 5.2 倍加速，在 CPU 上实现 8.6 倍加速，表明仅质心评分即可保持高召回率，并实现有效的早期过滤。
• 算法优化（质心交互与剪枝）与底层内核优化相结合，使 CPU 上的加速达到 42.4 倍（k=1000），而仅内核优化可实现 3 倍加速。
• 端到端延迟与数据集大小的平方根近似成比例，与 ColBERTv2 中质心数量相对于嵌入数量的缩放规律一致。
• 即使在 14000 万段落规模下，PLAID 在 CPU 上实现端到端延迟低于 100ms，在 GPU 上低于 50ms，支持大规模实时推理。
• PLAID 在 CPU 线程数增加时表现出良好扩展性，16 个线程下实现 4.9 倍加速，表明尽管负载均衡不够理想，仍具备有效的并行化能力。