[논문 리뷰] Learning Discrete Representations via Constrained Clustering for Effective and Efficient Dense Retrieval

Jingtao Zhan, Jiaxin Mao외 4명|arXiv (Cornell University)|2021. 10. 12.

Advanced Image and Video Retrieval Techniques인용 수 10

한 줄 요약

RepCONC는 제약된 클러스터링을 통해 이산 표현을 학습하는 밀도 검색의 새로운 방법을 제안한다. 이는 이중 인코더와 제품 양자화(PQ)를 동시에 최적화하여 효율적이고 종단 간(end-to-end) 학습이 가능하게 한다. 양자화를 균일한 클러스터링 조건을 가진 제약된 클러스터링 문제로 모델링함으로써, RepCONC는 뛰어난 검색 효과성, 메모리 효율성(색인 크기 <0.5GB), 그리고 CPU 호환성 속도를 달성하며, 즉각 검색 벤치마크에서 정확성과 효율성 면에서 최신 기준 모델들을 능가한다.

ABSTRACT

Dense Retrieval (DR) has achieved state-of-the-art first-stage ranking effectiveness. However, the efficiency of most existing DR models is limited by the large memory cost of storing dense vectors and the time-consuming nearest neighbor search (NNS) in vector space. Therefore, we present RepCONC, a novel retrieval model that learns discrete Representations via CONstrained Clustering. RepCONC jointly trains dual-encoders and the Product Quantization (PQ) method to learn discrete document representations and enables fast approximate NNS with compact indexes. It models quantization as a constrained clustering process, which requires the document embeddings to be uniformly clustered around the quantization centroids and supports end-to-end optimization of the quantization method and dual-encoders. We theoretically demonstrate the importance of the uniform clustering constraint in RepCONC and derive an efficient approximate solution for constrained clustering by reducing it to an instance of the optimal transport problem. Besides constrained clustering, RepCONC further adopts a vector-based inverted file system (IVF) to support highly efficient vector search on CPUs. Extensive experiments on two popular ad-hoc retrieval benchmarks show that RepCONC achieves better ranking effectiveness than competitive vector quantization baselines under different compression ratio settings. It also substantially outperforms a wide range of existing retrieval models in terms of retrieval effectiveness, memory efficiency, and time efficiency.

연구 동기 및 목표

• 밀도 검색 모델이 큰 밀도 벡터 색인과 완전한 근접 이웃 검색에 의존함에 따라 발생하는 높은 메모리 및 계산 비용을 해결하기 위해.
• 비가역적인 양자화 연산을 제약된 클러스터링을 통해 미분 가능하게 만들어 이중 인코더와 양자화 방법의 종단 간 공동 최적화를 가능하게 하기 위해.
• 균일한 클러스터링 조건을 통해 균형 잡히고 구분 가능한 이산 표현을 보장함으로써 검색 효과성과 효율성을 향상시키기 위해.
• 작은 색인을 갖는 벡터 기반의 역인verted 파일 시스템(IVF)을 사용하여 CPU에서 빠른 근사 근접 이웃 검색을 지원하기 위해.
• 동시에 검색 효과성, 메모리 효율성, 시간 효율성 면에서 최신 기술 수준의 성능을 달성하기 위해.

제안 방법

• RepCONC는 문서 임베딩을 PQ 중심점 주변으로 클러스터링하면서, 클러스터링 손실을 포함하는 제약된 클러스터링 과정으로 양자화를 모델링한다.
• 벡터가 소수의 클러스터에 집중되는 것을 방지하기 위해 균일한 클러스터링 조건을 도입하여, 양자화된 벡터가 유지되는 구분 가능한 표현을 보장한다.
• 효율적인 근사 최적화를 위해 이 조건을 최적 운반 문제로 완화한다.
• 역전파를 사용하여 이중 인코더와 PQ 중심점을 함께 학습함으로써, 두 구성 요소의 완전한 종단 간 최적화를 가능하게 한다.
• 지연 시간을 줄이기 위해 CPU에서 빠른 비완전한 벡터 검색을 가능하게 하는 벡터 기반의 역인verted 파일 시스템(IVF)을 활용한다.
• 추론 단계에서는 균일성 조건을 제거하여 효율성을 유지하면서도, 학습된 중심점과 할당을 통해 빠른 근사 근접 이웃 검색을 구현한다.

실험 결과

연구 질문

• RQ1 균일성 조건을 가진 제약된 클러스터링이 밀도 검색에서 이산 표현의 효과성 향상에 기여하는가?
• RQ2 양자화의 비가역성에도 불구하고 이중 인코더와 PQ의 종단 간 공동 최적화가 가능할 수 있는가?
• RQ3 RepCONC는 최신 기술 수준의 벡터 양자화 기반 모델 대비 더 뛰어난 검색 효과성, 메모리 효율성, 시간 효율성을 달성하는가?
• RQ4 RepCONC는 랭킹 성능을 희생시키지 않고도 효율적인 CPU 기반 검색을 가능하게 하는가?
• RQ5 균일한 클러스터링 조건은 클러스터 분포와 표현 품질에 어떤 영향을 미치는가?

주요 결과

• RepCONC는 48바이트당 패assing(BPP)에서 MRR@10이 0.340을 기록하여, 동일한 압축 비율에서 JPQ(0.332)와 DPQ(0.305)를 능가한다.
• 색인 크기를 0.5GB 이하로 줄여 COIL(60GB), ColBERT(162GB), uniCOIL(1.3GB), DeepImpact(1.5GB)보다 크게 작게 만들었다.
• RepCONC는 IVF를 활용하여 효율적인 CPU 기반 검색을 가능하게 하여, 고비용 GPU 가속을 요구하지 않고도 높은 성능을 달성한다.
• 제거 실험을 통해 균일한 클러스터링 조건이 필수적임을 확인: 이를 제거하면 클러스터 분포가 불균형해지고 성능 저하가 발생한다.
• 동적 하드 네거티브를 적용하면 RepCONC는 48 BPP에서 MRR@10을 0.340으로 향상시켜 공동 최적화의 효과를 입증한다.
• 제약된 클러스터링 접근법은 비제약 학습 대비 더 균형 잡힌 클러스터 할당을 생성하며, OPQ의 분포에 가까워지면서도 종단 간 미분 가능성을 유지한다.