A Library For Dense Retrieval Oriented Vector Quantization
Vector quantization (VQ) is widely applied to many ANN libraries, like FAISS, ScaNN, SPTAG, DiskANN to facilitate real-time and memory-efficient dense retrieval. However, conventional vector quantization methods, like IVF, PQ, OPQ, are not optimized for the retrieval quality. In this place, We present LibVQ, the first library developed for dense retrieval oriented vector quantization. LibVQ is highlighted for the following features:
-
Knowledge Distillation. The knowledge distillation based learning process can be directly applied to the off-the-shelf embeddings. It gives rise to the strongest retrieval performance in comparison with any existing VQ based ANN indexes.
-
Flexible usage and input conditions. LibVQ may flexibly support different usages, e.g., training VQ parameters only, or joint adaptation of query encoder. LibVQ is designed to handle a wide range of input conditions: it may work only with off-the-shelf embeddings; it may also leverage extra data, e.g., relevance labels, and source queries, for further enhancement.
-
Learning and Deployment. The learning is backended by PyTorch, which can be easily configured for the efficient training based on different computation resources. The well-trained VQ parameters are wrapped up with FAISS backend ANN indexes, e.g., IndexPQ, IndexIVFPQ, etc., which are directly deployable for large-scale dense retrieval applications.
- From source
git clone https://github.com/staoxiao/LibVQ.git
cd LibVQ
pip install .
In LibVQ, users can construct a index and train it by a simple way. Please refer to our docs for more details. Besides, we provide some examples below to illustrate the usage of LibVQ.
- IVFPQ (Compression Ratio = 96)
Methods | MRR@10 | Recall@10 | Recall@100 |
---|---|---|---|
Faiss-IVFPQ | 0.1380 | 0.2820 | 0.5617 |
Faiss-IVFOPQ | 0.3102 | 0.5593 | 0.8148 |
Scann | 0.1791 | 0.3499 | 0.6345 |
LibVQ(contrastive_index) | 0.3179 | 0.5724 | 0.8214 |
LibVQ(distill_index) | 0.3253 | 0.5765 | 0.8256 |
LibVQ(distill_index_nolabel) | 0.3234 | 0.5813 | 0.8269 |
LibVQ(contrastive_index-and-query-encoder) | 0.3192 | 0.5799 | 0.8427 |
LibVQ(distill_index-and-query-encoder) | 0.3311 | 0.5907 | 0.8429 |
LibVQ(distill_index-and-query-encoder_nolabel) | 0.3285 | 0.5875 | 0.8401 |
- PQ (Compression Ratio = 96)
Methods | MRR@10 | Recall@10 | Recall@100 |
---|---|---|---|
Faiss-PQ | 0.1145 | 0.2369 | 0.5046 |
Faiss-OPQ | 0.3268 | 0.5939 | 0.8651 |
Scann | 0.1795 | 0.3516 | 0.6409 |
LibVQ(distill_index) | 0.3435 | 0.6203 | 0.8825 |
LibVQ(distill_index_nolabel) | 0.3467 | 0.6180 | 0.8849 |
LibVQ(distill_index-and-query-encoder) | 0.3446 | 0.6201 | 0.8837 |
LibVQ(distill_index-and-two-encoders) | 0.3475 | 0.6223 | 0.8901 |
- IVFPQ (Compression Ratio = 384)
Methods | Recall@5 | Recall@10 | Recall@20 | Recall@100 |
---|---|---|---|---|
Faiss-IVFPQ | 0.1504 | 0.2052 | 0.2722 | 0.4523 |
Faiss-IVFOPQ | 0.3332 | 0.4279 | 0.5110 | 0.6817 |
Scann | 0.2526 | 0.3351 | 0.4144 | 0.6016 |
LibVQ(contrastive_index) | 0.3398 | 0.4415 | 0.5232 | 0.6911 |
LibVQ(distill_index) | 0.3952 | 0.4900 | 0.5667 | 0.7232 |
LibVQ(distill_index_nolabel) | 0.4066 | 0.4936 | 0.5759 | 0.7301 |
LibVQ(contrastive_index-and-query-encoder) | 0.3548 | 0.4470 | 0.5390 | 0.7120 |
LibVQ(distill_index-and-query-encoder) | 0.4725 | 0.5681 | 0.6429 | 0.7739 |
LibVQ(distill_index-and-query-encoder_nolabel) | 0.4977 | 0.5822 | 0.6484 | 0.7764 |
- PQ (Compression Ratio = 384)
Methods | Recall@5 | Recall@10 | Recall@20 | Recall@100 |
---|---|---|---|---|
Faiss-PQ | 0.1301 | 0.1861 | 0.2495 | 0.4188 |
Faiss-OPQ | 0.3166 | 0.4105 | 0.4961 | 0.6836 |
Scann | 0.2526 | 0.3351 | 0.4144 | 0.6013 |
LibVQ(distill_index) | 0.3817 | 0.4806 | 0.5681 | 0.7357 |
LibVQ(distill_index_nolabel) | 0.3880 | 0.4858 | 0.5819 | 0.7423 |
LibVQ(distill_index-and-query-encoder) | 0.4709 | 0.5689 | 0.6481 | 0.7930 |
LibVQ(distill_index-and-query-encoder_nolabel) | 0.4883 | 0.5903 | 0.6678 | 0.7914 |
LibVQ(distill_index-and-two-encoders) | 0.5637 | 0.6515 | 0.7171 | 0.8257 |
LibVQ(distill_index-and-two-encoders_nolabel) | 0.5285 | 0.6144 | 0.7296 | 0.8096 |
-
Distii-VQ: Unifies the learning of IVF and PQ within a knowledge distillation framework. Accpted as a full paper by SIGIR 2022.
-
BiDR: Applies the learnable PQ in large-scale index and proposes the progressively optimized docs' embeddings for the better retrieval performance. Accpted as a full paper by WWW 2022.
-
MoPQ: This work identifies the limitation of using reconstruction loss minimization as the training objective of learnable PQ and proposes the Multinoulli Contrastive Loss. Accpted as a full paper by EMNLP 2021.