[论文解读] Distill-VQ: Learning Retrieval Oriented Vector Quantization By Distilling Knowledge from Dense Embeddings

Shitao Xiao, Zheng Liu+11 more|arXiv (Cornell University)|Apr 1, 2022

Advanced Image and Video Retrieval Techniques被引用 4

一句话总结

Distill-VQ 提出了一种知识蒸馏框架，通过模仿密集嵌入模型的 relevance 预测来训练向量量化（VQ）模块——IVF 和 PQ，以保留检索性能，从而在无需标注查询的情况下，在 MS MARCO 和 Natural Questions 数据集上实现了最先进的召回率和 MRR，即使在高压缩和低延迟条件下也表现优异。

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 组件，以在高压缩和低延迟条件下提升检索质量。
• 证明在检索系统中，高质量向量量化不再需要标注数据。

提出的方法

• 使用预训练的密集嵌入作为“教师”，预测查询与候选文档之间的相关性得分。
• 将 VQ 模块（IVF 和 PQ）视为“学生”，通过训练使其复现教师的相关性预测。
• 应用 ListNet 损失以确保教师和学生预测的相关性得分保持排序一致性。
• 从 Top-K 和同批候选文档中采样，构建知识蒸馏信号，而无需真实标签。
• 在统一的蒸馏框架内联合优化 IVF 和 PQ 组件，以保持检索性能。
• 将训练好的 VQ 模块集成到 FAISS 的 IVFPQ 索引中，以实现高效推理与评估。

实验结果

研究问题

• RQ1 能否通过从密集嵌入中蒸馏知识，在不依赖标注查询-文档对的情况下有效提升向量量化性能？
• RQ2 与直接最小化重建损失或检索损失相比，从密集嵌入中蒸馏相关性预测在 VQ 训练中表现如何？
• RQ3 在高斯压缩和低延迟条件下，Distill-VQ 在召回率和 MRR 上相比现有 VQ 方法能提升多少？
• RQ4 蒸馏对 IVF 和 PQ 组件分别产生何种影响，以及其如何影响效率-准确率权衡？
• RQ5 当标注数据稀缺或不可用时，Distill-VQ 是否仍能实现最先进性能？

主要发现

• 在 MS MARCO 和 Natural Questions 基准测试中，Distill-VQ 在召回率和 MRR 上显著优于 SOTA 的 VQ 方法。
• 在低比特率（例如，192 倍压缩）下，Distill-VQ 在 Recall@100 和 MRR@10 上仍保持显著优势，尤其在高斯压缩和低延迟条件下。
• 与 FAISS 原始 IVFPQ 相比，该方法在小码书尺寸下实现了更优的效率-召回权衡。
• 即使在小码书（例如，8×8）条件下，Distill-VQ 仍优于使用更大码书的基线模型，表明其具备强大的泛化能力和数据效率。
• 采用 Top-K + 同批采样策略，可在无需标注数据的情况下实现有效蒸馏，使方法具备可扩展性和实用性。
• Distill-VQ 在不同压缩比率和搜索设置下均保持一致的性能增益，展现出强鲁棒性和泛化能力。