[논문 리뷰] Distill-VQ: Learning Retrieval Oriented Vector Quantization By Distilling Knowledge from Dense Embeddings

Shitao Xiao, Zheng Liu외 11명|arXiv (Cornell University)|2022. 04. 01.

Advanced Image and Video Retrieval Techniques인용 수 4

한 줄 요약

Distill-VQ는 밀도 임베딩에서 유사도 예측을 모방함으로써 검색 성능을 유지하는 지식 distillation 프레임워크를 제안한다. 이는 IVF와 PQ를 포함한 벡터 양자화(VQ) 모듈을 훈련시키며, 레이블이 없는 쿼리 데이터를 요구하지 않아도 MS MARCO와 Natural Questions에서 최신 기준 성능을 달성한다. 고압축 및 저지연 조건에서도 성능이 뛰어나다.

ABSTRACT

Vector quantization (VQ) based ANN indexes, such as Inverted File System (IVF) and Product Quantization (PQ), have been widely applied to embedding based document retrieval thanks to the competitive time and memory efficiency. Originally, VQ is learned to minimize the reconstruction loss, i.e., the distortions between the original dense embeddings and the reconstructed embeddings after quantization. Unfortunately, such an objective is inconsistent with the goal of selecting ground-truth documents for the input query, which may cause severe loss of retrieval quality. Recent works identify such a defect, and propose to minimize the retrieval loss through contrastive learning. However, these methods intensively rely on queries with ground-truth documents, whose performance is limited by the insufficiency of labeled data. In this paper, we propose Distill-VQ, which unifies the learning of IVF and PQ within a knowledge distillation framework. In Distill-VQ, the dense embeddings are leveraged as "teachers", which predict the query's relevance to the sampled documents. The VQ modules are treated as the "students", which are learned to reproduce the predicted relevance, such that the reconstructed embeddings may fully preserve the retrieval result of the dense embeddings. By doing so, Distill-VQ is able to derive substantial training signals from the massive unlabeled data, which significantly contributes to the retrieval quality. We perform comprehensive explorations for the optimal conduct of knowledge distillation, which may provide useful insights for the learning of VQ based ANN index. We also experimentally show that the labeled data is no longer a necessity for high-quality vector quantization, which indicates Distill-VQ's strong applicability in practice.

연구 동기 및 목표

• 기존의 벡터 양자화(VQ) 목표(재구성 오차 최소화)와 실제 검색 성능 사이의 괴리 문제를 해결한다.
• 기존의 검색 최적화 VQ 방법이 대량의 비용이 들고 제한된 레이블이 있는 쿼리-문서 쌍을 기반으로 한 대비 학습에 의존하는 문제를 해결한다.
• 밀도 임베딩을 지식 교사로 활용하여 막대한 레이블이 없는 데이터를 통해 VQ 모듈을 효과적으로 훈련시킬 수 있도록 한다.
• IVF와 PQ 컴포넌트를 함께 distillation을 통해 최적화하여 고압축 및 저지연 조건에서도 검색 품질을 향상시킨다.
• 레이블이 없는 데이터로도 고성능의 벡터 양자화를 달성할 수 있음을 입증한다.

제안 방법

• 사전 훈련된 밀도 임베딩을 '교사'로 사용하여 쿼리와 후보 문서 간의 관련성 점수를 예측한다.
• VQ 모듈(IVF와 PQ)을 '학생'으로 간주하고, 교사의 관련성 예측을 재현하도록 훈련시킨다.
• 교사와 학생의 예측 관련성 점수 간 순서 유지 성질을 보장하기 위해 ListNet 손실을 적용한다.
• 지표 레이블이 필요 없이 Top-K 및 인라인 배치 내 후보 문서를 샘플링하여 지식 distillation 신호를 구성한다.
• 통합된 distillation 프레임워크 내에서 IVF와 PQ 컴포넌트를 함께 최적화하여 검색 성능을 유지한다.
• 훈련된 VQ 모듈을 FAISS의 IVFPQ 인덱스에 통합하여 효율적인 추론과 평가를 수행한다.

실험 결과

연구 질문

• RQ1 레이블이 없는 쿼리-문서 쌍에 의존하지 않고도 밀도 임베딩에서 유지를 통해 VQ 성능을 효과적으로 향상시킬 수 있는가?
• RQ2 VQ 훈련에서 재구성 오차나 검색 손실을 직접 최소화하는 것과 비교해, 밀도 임베딩에서의 관련성 예측을 distillation하는 방식은 어떠한가?
• RQ3 고압축 및 저지연 조건에서 Distill-VQ는 기존 VQ 방법에 비해 Recall과 MRR 측면에서 어느 정도 뛰어나게 성능을 발휘하는가?
• RQ4 Distillation은 IVF와 PQ 컴포넌트에 각각 어떤 영향을 미치며, 이는 효율성-정확도 트레이드오프에 어떤 영향을 미치는가?
• RQ5 레이블 데이터가 부족하거나 없을 경우에도 Distill-VQ는 최신 기준 성능을 달성할 수 있는가?

주요 결과

• Distill-VQ는 MS MARCO와 Natural Questions 벤치마크에서 모두 SOTA VQ 방법보다 Recall과 MRR 측면에서 뚜렷한 승리를 거두었다.
• 낮은 비트레이트(예: 192× 압축) 조건에서도 Distill-VQ는 Recall@100과 MRR@10에서 상당한 우위를 보였으며, 특히 고압축 및 저지연 환경에서 두드러졌다.
• 코드북 크기가 작을 경우 기존 FAISS의 원본 IVFPQ보다 더 나은 효율-정확도 트레이드오프를 달성했다.
• 작은 코드북(예: 8×8) 조건에서도 더 큰 코드북을 사용한 베이스라인을 초월하는 성능을 기록하여 강력한 일반화 능력과 데이터 효율성을 입증했다.
• Top-K + 인라인 배치 샘플링을 활용함으로써 레이블이 없는 데이터로도 효과적인 distillation을 가능하게 하여, 이 방법이 확장 가능하고 실용적임을 입증했다.
• 다양한 압축 비율과 검색 설정에서도 일관된 성능 향상을 유지하여 강건성과 일반화 능력을 입증했다.