[논문 리뷰] PLAID: An Efficient Engine for Late Interaction Retrieval

Keshav Santhanam, Omar Khattab외 2명|arXiv (Cornell University)|2022. 05. 19.

Topic Modeling인용 수 8

한 줄 요약

PLAID는 ColBERTv2를 사용하여 후기 상호작용 검색을 가속화하는 고도로 최적화된 검색 엔진으로, 중심점 상호작용과 중심점 프루닝을 도입하여 GPU에서는 최대 7배, CPU에서는 최대 45배의 검색 지연 감소를 이끌어내며, 최대 140M개의 패assage를 포함한 대규모 데이터셋에서 최신 기술 수준의 검색 품질을 유지한다.

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억 개 이상의 패스제이지가 포함된 생산 환경에서 100M+ 패스제이지에서 CPU와 GPU 모두에서 저지연 추론을 가능하게 한다.
• 후보 스코링 중 계산 및 메모리 오버헤드를 크게 줄이면서도 검색 품질을 유지한다.
• 중심점 ID의 이산성에 기반해 효율적인 필터링과 스코링을 가능하게 하는 모듈형이며 고도로 최적화된 엔진을 개발한다.

제안 방법

• 중심점 상호작용 도입: 각 패스제이지를 정수 중심점 ID의 경량 버전으로 간주하여, 전체 잔류 벡터를 로드하지 않고도 빠르고 근사적인 스코링을 가능하게 한다.
• 사전 계산된 중심점-질의 거리를 사용하여 검색 중에 모든 중심점에 대한 유사도 스코어를 재계산하지 않도록 하여, 여러 패스제이지 간에 재사용이 가능하도록 한다.
• 중심점 프루닝 구현: 낮은 스코어를 가진 중심점의 희소성에 기반해 초기 필터링 단계에서 먼 중심점 ID를 건너뛰어 후보 세트 크기를 줄인다.
• GPU용 팯팅 없는 MaxSim 커널과 CPU용 최적화된 C++ 커널을 설계하여 잔류 복구 및 스코링 연산을 가속화한다.
• 완전한 잔류 스코링 이전에 중심점 상호작용과 프루닝을 적용하는 다단계 검색 파이프라인을 설계하여, 저품질 후보에 대한 비용이 많이 드는 연산을 최소화한다.
• 중심점의 고정된 이산 어휘를 활용해 CPU 및 GPU 실행에서 효율적인 벡터화된 연산과 메모리 액세스 패턴을 가능하게 한다.

실험 결과

연구 질문

• RQ1 중심점 기반 표현만으로도 후기 상호작용 검색에서 높은 리콜을 달성할 수 있는가? 이는 전체 잔류 스코링 이전의 효과적인 필터링을 가능하게 하는가?
• RQ2 중심점 프루닝이 검색 품질을 떨어뜨리지 않으면서 후보 패스제이지 수를 얼마나 효과적으로 줄일 수 있는가?
• RQ3 알고리즘 최적화(중심점 상호작용 및 프루닝)와 저수준 커널 최적화를 조합했을 때 CPU 및 GPU에서 엔드 투 엔드 지연을 얼마나 향상시킬 수 있는가?
• RQ4 다양한 하드웨어 구성에서 코퍼스 크기와 검색 깊이(k)가 증가함에 따라 PLAID는 어떻게 스케일링되는가?
• RQ5 초기 단계에서 전체 잔류 스코링을 중심점 전용 스코링으로 대체할 경우, 지연 감소와 검색 품질 간의 상호 교환 관계는 어떠한가?

주요 결과

• PLAID ColBERTv2는 MS MARCO v1, v2, Wikipedia, LoTTE 벤치마크에서 바닐라 ColBERTv2 대비 GPU에서 최대 7배, CPU에서 최대 45배의 속도 향상을 기록했으며, 품질 손실이 측정되지 않을 정도로 유지된다.
• 중심점 상호작용만으로도 GPU에서 5.2배, CPU에서 8.6배의 속도 향상이 이루어졌으며, 중심점 전용 스코링이 높은 리콜을 유지하고 초기 필터링에 효과적임을 입증한다.
• 알고리즘 최적화(중심점 상호작용 및 프루닝)와 저수준 커널 최적화를 조합하면 k=1000일 때 CPU에서 42.4배의 속도 향상을 기록했으며, 커널 최적화만으로는 3배에 그친다.
• 엔드 투 엔드 지연은 데이터셋 크기의 제곱근에 비례하여 스케일링되며, 이는 ColBERTv2에서 임베딩 수 대비 중심점 수의 스케일링과 일치한다.
• 140M개 패스제이지에서도 CPU에서는 100ms 미만, GPU에서는 50ms 이내의 엔드 투 엔드 지연을 달성하여 대규모 실시간 추론을 가능하게 한다.
• CPU 스레드 수 증가에 따라 잘 스케일링되며, k=1000일 때 16개 스레드로 4.9배의 속도 향상을 기록했으며, 비효율적인 로드 밸런싱에도 불구하고 효과적인 병렬 처리가 가능함을 보여준다.