[论文解读] Learning Discrete Representations via Constrained Clustering for Effective and Efficient Dense Retrieval

Jingtao Zhan, Jiaxin Mao+4 more|arXiv (Cornell University)|Oct 12, 2021

Advanced Image and Video Retrieval Techniques被引用 10

一句话总结

RepCONC 提出了一种用于密集检索的新方法，通过约束聚类学习离散表示，联合优化双编码器和乘积量化（PQ），实现高效、端到端的训练。通过将量化建模为具有均匀聚类约束的约束聚类问题，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.

研究动机与目标

• 为解决密集检索模型的高内存和计算成本问题，这些模型依赖于大型密集向量索引和全量最近邻搜索。
• 通过将不可微的量化操作通过约束聚类变为可微，实现双编码器和量化方法的端到端训练。
• 通过均匀聚类约束确保聚类分布均衡、可区分，从而提升检索效果和效率。
• 通过基于向量的倒排文件系统（IVF）实现紧凑索引，在 CPU 上支持快速的近似最近邻搜索。
• 同时实现最先进的检索效果、内存效率和时间效率。

提出的方法

• RepCONC 将量化建模为约束聚类过程，其中文档嵌入围绕 PQ 中心点聚类，并通过聚类损失进行优化。
• 引入均匀聚类约束以防止向量集中在少数聚类中，确保量化后的向量保持可区分性。
• 将该约束松弛为最优传输问题，以在端到端训练期间实现高效的近似优化。
• 通过反向传播联合训练双编码器和 PQ 中心点，实现对两个组件的完整端到端优化。
• 采用基于向量的倒排文件系统（IVF）在 CPU 上实现快速、非全量的向量搜索，降低延迟。
• 推理过程中，丢弃均匀约束以保持效率，而学习到的中心点和分配关系则支持快速的近似最近邻搜索。

实验结果

研究问题

• RQ1 具有均匀性约束的约束聚类能否提升密集检索中离散表示的效果？
• RQ2 尽管量化操作不可微，能否实现双编码器和 PQ 的端到端联合优化？
• RQ3 与最先进向量量化基线相比，RepCONC 是否在检索效果、内存效率和时间效率方面表现更优？
• RQ4 RepCONC 能否在不牺牲排序性能的前提下，实现高效的基于 CPU 的检索？
• RQ5 均匀聚类约束如何影响聚类分布和表示质量？

主要发现

• 在每段 48 字节（BPP）的压缩比下，RepCONC 的 MRR@10 达到 0.340，优于 JPQ（0.332）和 DPQ（0.305）。
• 模型将索引大小减少至 0.5GB 以下，显著小于 COIL（60GB）、ColBERT（162GB）、uniCOIL（1.3GB）和 DeepImpact（1.5GB）。
• RepCONC 支持高效的基于 CPU 的检索，通过 IVF 实现高性能，无需昂贵的 GPU 加速。
• 消融实验证实，均匀聚类约束至关重要：移除该约束后，聚类分布变得极不平衡，性能显著下降。
• 在动态难样本训练下，RepCONC 进一步将 MRR@10 提升至 48 BPP 时的 0.340，证明了联合优化的有效性。
• 约束聚类方法产生的聚类分配比无约束训练更均衡，接近 OPQ 的分布，同时保持了端到端可微性。