Learning Discrete Representations via Constrained Clustering for Effective and Efficient Dense Retrieval (Extended Abstract)

1 Introduction

Dense Retrieval (DR) has become a popular paradigm for first-stage retrieval in ad-hoc retrieval tasks. It embeds queries and documents in a latent vector space with dualencoders and uses nearest neighbor search to retrieve relevant documents. With end-to-end supervised training, DR models have achieved state-of-the-art ranking performance and significantly outperform BoW models [Zhan et al., 2021b; Lin et al., 2020; Xiong et al., 2021].

Despite the success in improving ranking performance, most existing DR models [Zhan et al., 2021b; Xiong et al., 2021; Karpukhin et al., 2020] are inefficient in memory usage and retrieval speed. For memory inefficiency, the size of the embedding index is usually an order of magnitude larger than that of BoW index [Zhan et al., 2021a]. At runtime, the vectors must be loaded to costly system memory or even GPU memory. As for time inefficiency, many existing DR models conduct exhaustive search, i.e., computing relevance scores between the submitted query and all documents. As a result, these DR models cannot use CPUs for retrieval due to high latency and have to use much more expensive GPUs to accelerate the search.

∗This is an abridged version of the paper [Zhan et al., 2022] that won the Best Paper Award at ACM WSDM 2022.

†Corresponding Author

To tackle this problem, we present RepCONC, which stands for learning discrete Representations via CONstrained Clustering 1 . RepCONC learns discrete representations with Product Quantization (PQ) so that the representations can be encoded into compact indexes for efficient vector search. RepCONC utilizes joint optimization of dual-encoders and PQ to achieve effective ranking results. During joint optimization, RepCONC models quantization as a constrained clustering process, which involves a clustering loss and a uniform clustering constraint. The clustering loss is introduced to train the discrete codes. And the uniform clustering constraint facilitates distinguishability of discrete representations by requiring the vectors to be equally assigned to all quantization centroids. Besides constrained clustering, RepCONC further employs vector-based inverted file system (IVF) [Jegou et al., 2010] to enable efficient non-exhaustive vector search on either GPU or CPU.

We conduct experiments on two widely-adopted ad-hoc retrieval benchmarks [Bajaj et al., 2016; Craswell et al., 2020] and compare RepCONC with a wide range of baselines. Experimental results show that: 1) RepCONC significantly outperforms competitive vector compression baselines with different compression ratio settings. 2) RepCONC substantially outperforms various retrieval baselines in terms of retrieval effectiveness, memory efficiency, and time efficiency.

2 Constrained ClusteringModel

In this section, we propose RepCONC, which stands for learning discrete Representations via CONstrained Clustering. We first introduce the preliminary of Production Quantization [Jegou et al., 2010] and then elaborate on our model.

1 Code and models are available at https://github.com/ jingtaozhan/RepCONC.

2.1 Revisiting Product Quantization

RepCONC is based on Product Quantization (PQ) [Jegou et al., 2010]. For vectors of dimension D, PQ defines M sets of embeddings, each of which includes K embeddings of dimension D/M. They are called PQ Centroid Embeddings. Formally, let ci,j be the jth centroid embedding from the ith set:

Given a document embedding d ∈ RD , PQ firstly splits it into M sub-vectors.

Then PQ independently quantizes each sub-vector to the nearest PQ Centroid Embedding. Formally, to quantize a subvector di, PQ selects the nearest ci,ϕi(d):

Let ϕ(d) be the concatenation of ϕi(d):

where comma denotes vector concatenation. ϕ(d) is called the Index Assignment of d. ϕ(d) can reconstruct the quantized document embedding ˆ d as follows:

PQ does not explicitly store d or ˆ d. It only stores the PQ Centroid Embeddings {ci,j} and Index Assignments ϕ(d).

2.2 Clustering and Representation Learning

RepCONC views joint optimization as a simultaneous clustering and representation learning problem. It utilizes both the ranking-oriented loss [Zhan et al., 2021a] and a clustering loss. We illustrate the training workflow in Figure 1.

The ranking-oriented loss computes the ranking loss based on the compressed document embeddings. Therefore, it better evaluates the ranking performance with respect to the current compression parameters. Let d+ and d- be relevant and irrelevant documents, respectively. Ranking-oriented loss Lr is formulated as:

RepCONC regards quantization as a clustering problem and additionally introduces the MSE loss Lm:

Minimizing Lm requires the document embeddings to be clustered around the centroid embeddings. Combining both Lr and Lm helps the model to cluster document embeddings based on ranking effectiveness. It is expected to produce better clustering compared with unsupervised training. The final loss L is a weighted sum of ranking-oriented loss Lr and the MSE loss Lm.

Since quantization is non-differentiable, we explicitly design the gradient back-propagation policy for document encoders. The gradients of uncompressed document embeddings are defined as follows:

As the equation shows, we add the gradient of quantized document embeddings (the first term). The gradients are further back-propagated to document encoders. Gradients of other parameters can be derived with chain rule.

2.3 Importance of UniformClustering

It is non-trivial to simultaneously conduct clustering and representation learning because the two objects are conflicting to some extent. Although representation learning encourages vectors to be distinguishable, clustering encourages vectors to be identical. In the iterative training process, clustering objective leads to unbalanced clustering distribution, which affects the vector distinguishability and compromises ranking effectiveness.

We tackle this challenge by imposing a uniform clustering constraint. It requires the document sub-vectors to be equally assigned to all PQ Centroid Embeddings. The learning object along with the constraint is formally expressed as:

We illustrate constrained clustering in Figure 2. As the figure shows, the discrete document embeddings are selected by minimizing the quantization error (maximizing the similarity) given the uniform clustering constraint. Without the constraint, the discrete document embeddings become identical. Now we theoretically analyze the importance of uniform clustering. Due to the limitation of space, we summarize the theoretical conclusions here and refer readers to the full paper [Zhan et al., 2022] for detailed derivation.

Theorem 1. Maximizing the distinguishability of vectors is equivalent to forcing the vectors to be equally quantized to all possible Index Assignments.

That is to say, quantizing vectors equally to all possible Index Assignments helps representations to be distinguishable.

Theorem 2. Uniformly clustering sub-vectors is the essential condition of maximum distinguishability.

Theorem 3. If sub-vectors are independent, uniformly clustering sub-vectors is the sufficient condition of maximum distinguishability.

Although independence among sub-vectors may not hold for practical dual-encoders, we believe uniformly clustering sub-vectors is still helpful for distinguishing quantized vectors.

2.4 Constrained Clustering Optimization

This section shows how to incorporate the uniform clustering constraint to select Index Assignments during training.

We introduce a posterior distribution q(j|di), which is the probability that the sub-vector di is quantized to the centroid ci,j. The Index Assignment, ϕi(d), is the centroid with the maximum probability:

For PQ that uses Eq. (3), q(j|di) can be regarded as being computed solely based on quantization error. Here for RepCONC, we compute q(j|di) by minimizing the quantization error given the uniform clustering constraint:

where D indicates the set of all documents. The first condition constrains q(j|di) to be binary, the second condition is a natural requirement for probability, and the third condition is exactly the uniform clustering constraint. Without the third condition, Eq. (11) and (12) degenerate to Eq. (3), i.e., selecting Index Assignments with minimum quantization error.

Solving Eq. (12) is particularly difficult because it is a combinatorial optimization problem with the scale of millions or even billions of documents. Therefore, we use an approximate solution by relaxing q to be continuous and focusing on uniformly clustering a mini-batch of documents B:

Since ‖ci,j -di‖2 can be regarded as the cost of mapping di to ci,j , this is an instance of the optimal transport problem and can be solved in polynomial time by linear program. In our implementation, we use Sinkhorn-Knopp algorithm [Cuturi, 2013] to efficiently solve Eq. (13).

2.5 Accelerating Search with IVF

Besides PQ, RepCONC employs the inverted file system (IVF) to accelerate vector search. After quantizing document embeddings, RepCONC uses k-means to generate n clusters. Given a query embedding, RepCONC selects the nearest ˜n clusters and only ranks the documents in them. The documents in other clusters are ignored. In this way, RepCONC approximately accelerates vector search by n/˜n.

3 Experimental Setup

Here we present our experimental settings.

3.1 Datasets andMetrics

We conduct experiments on two large-scale ad-hoc retrieval benchmarks from the TREC 2019 Deep Learning Track [Craswell et al., 2020; Bajaj et al., 2016], passage ranking and document ranking. Due to the limited space, this paper only reports the performance on the passage ranking task. Please refer to our full paper [Zhan et al., 2022] for comprehensive results. The passage ranking task has a corpus of 8.8M passages, 0.5M training queries, 7k development queries (henceforth, MARCO Passage), and 43 test queries (DL Passage). We report the official metrics and R@100 based on the full-corpus retrieval results.

3.2 Baselines

We exploit two types of baselines, vector compression methods and retrieval models.

For vector compression methods, we adopt both unsupervised and supervised methods. The former include PQ [Jegou et al., 2010], ScaNN [Guo et al., 2020], ITQ+LSH [Gong et al., 2012], OPQ [Ge et al., 2013], and OPQ+ScaNN. The latter include DPQ [Chen et al., 2020] and JPQ [Zhan et al., 2021a].

For retrieval baselines, we utilize BoW models, dualencoders, and some competitive complex retrieval systems. BoW models involve BM25 [Robertson and Walker, 1994], DeepCT [Dai and Callan, 2019], doc2query [Nogueira et al., 2019b], and docT5query [Nogueira et al., 2019a]. Dual-encoders include RepBERT [Zhan et al., 2020], ANCE [Xiong et al., 2021], and ADORE [Zhan et al., 2021b]. Complex retrieval systems are much slower than BoW and dual-encoders. They include ColBERT [Khattab and Zaharia, 2020], COIL [Gao et al., 2021], uniCOIL [Lin and Ma, 2021], and DeepImpact [Mallia et al., 2021].

4 Experiments

We empirically evaluate RepCONC in this section.

4.1 Comparison with CompressionMethods

This section compares RepCONC with vector compression baselines to answer RQ1. Ranking performances in terms of different compression ratios are plotted in Figure 3. The advantage of RepCONC is more significant when larger compression ratios are used. For example, its MRR score is more than twice the JPQ's score when the compression ratio is 784x. We believe this is because RepCONC is able to generate high-quality Index Assignments specifically for ranking effectiveness, which becomes more important when fewer bytes are used. Instead, JPQ uses K-Means to produce task-blind Index Assignments and compromises ranking performance.

4.2 Comparison with Retrieval Models

This section compares RepCONC with various retrieval models to address RQ2. We firstly compare it with efficient firststage retrievers and then compare it with complex (slow) endto-end retrievers.

Comparison with First-Stage Retrievers

Figure 4 summarizes the effectiveness-memory tradeoff. As the figure shows, although DR models are much more effective than BoW models, they incur severe memory inefficiency. By jointly training the dual-encoders and quantization methods, RepCONC substantially improves memory efficiency of DR while still being very effective in ranking. It outperforms RepBERT [Zhan et al., 2020] and ANCE [Xiong et al., 2021] in effectiveness, and is almost as effective as ADORE [Zhan et al., 2021b], the state-of-the-art DR model trained by negative sampling.

Comparison with Complex End-to-End Retrievers

This section compares RepCONC with some complex (slow) end-to-end neural retrieval models. These models achieve better ranking performance with much higher query latency because of their complex model architecture. In consideration of fair comparison, we add a reranking stage to RepCONC and compare them in terms of effectiveness-latency tradeoff. The reranking models are MonoBERT and DuoT5 models open-sourced by the pygaggle library 2 . Note, query encoding and reranking are performed on GPU while the search is performed on CPU with one thread. Ranking results are shown in Figure 5. We can see that RepCONC-IVF+Rerank substantially outperforms all baselines in terms of both effectiveness and time efficiency.

5 Conclusions

To solve the efficiency issue existing in brute-force DR models, we present RepCONC, which learns discrete representations by modeling quantization as constrained clustering in the joint learning process. The clustering object requires the document embeddings to be clustered around the quantization centroids and facilitates joint optimization of PQ parameters and dual-encoders. We also introduce a uniform clustering constraint to maximize the representation distinguishability. We conduct experiments on widely-adopted ad-hoc retrieval benchmarks. Experimental results show that RepCONC significantly outperforms competitive quantization baselines and substantially improves the memory efficiency and time efficiency of DR.

Acknowledgments

This work is supported by the National Key Research and Development Program of China (2018YFC0831700), Natural Science Foundation of China (Grant No. 61732008, 61532011, 61902209, U2001212), Beijing Academy of Artificial Intelligence (BAAI), Tsinghua University Guoqiang Research Institute, Beijing Outstanding Young Scientist Program (NO. BJJWZYJH012019100020098) and Intelligent Social Governance Platform, Major Innovation & Planning Interdisciplinary Platform for the "Double-First Class" Initiative, Renmin University of China.

2 https://github.com/castorini/pygaggle