Title: Knowledge Distillation for High Dimensional Search Index
Venue: NeurIPS 2023
Accepted: yes
Keywords: Approximate Nearest Neighbor Search

Knowledge Distillation

Product Quantization

Inverted Index
Abstract: Lightweight compressed models are prevalent in Approximate Nearest Neighbor Search (ANNS) and Maximum Inner Product Search (MIPS) owing to their superiority of retrieval efficiency in large-scale datasets. However, results given by compressed methods are less accurate due to the curse of dimension and the limitations of optimization objectives (e.g., lacking interactions between queries and documents). Thus, we are encouraged to design a new learning algorithm for the compressed search index on high dimensions to improve retrieval performance. In this paper, we propose a novel KnowledgeDistillation for high dimensional search index framework (KDindex), with the aim of efficiently learning lightweight indexes by distilling knowledge from high-precision ANNS and MIPS models such as graph-based indexes. Specifically, the student is guided to keep the same ranking order of the top-k relevant results yielded by the teacher model, which acts as the additional supervision signals between queries and documents to learn the similarities between documents. Furthermore, to avoid the trivial solutions that all candidates are partitioned to the same centroid, the reconstruction loss that minimizes the compressed error, and the posting list balance strategy that equally allocates the candidates, are integrated into the learning objective. Experiment results demonstrate that KDindex outperforms existing learnable quantization-based indexes and is 40× lighter than the state-of-the-art non-exhaustive methods while achieving comparable recall quality.

Knowledge Distillation for High Dimensional Search Index

Abstract

Lightweight compressed indexes are prevalent in Approximate Nearest Neighbor Search (ANNS) and Maximum Inner Product Search (MIPS) owing to their superiority of retrieval efficiency in large-scale datasets. However, results given by compressed indexes are less accurate due to the curse of dimension and limitation of optimization objectives (e.g., lacking interactions between queries and documents). Thus, we are encouraged to design a new learning algorithm for the compressed search index in high dimensions to improve retrieval performance. In this paper, we propose a novel K nowledge D istillation for high dimensional search index framework ( KDindex ), with the aim of efficiently learning lightweight indexes by distilling knowledge from high-precision ANNS and MIPS models such as graph-based indexes. Specifically, the student is guided to keep the same ranking order of the top-k relevant results yielded by the teacher model, which acts as the additional supervision signals between queries and documents to learn the similarities between documents. Furthermore, to avoid the trivial solutions that all candidates are partitioned to the same post list, the reconstruction loss that minimizes the compressed error, and the posting list balance strategy that equally allocates the candidates, are integrated into the learning objective. Experiment results demonstrate that KDindex outperforms existing learnable quantization-based indexes and is 40× lighter than the state-of-the-art non-exhaustive methods while achieving comparable recall quality.

1 Introduction

Vector nearest neighbor search, which retrieves the most relevant vectors with the maximum similarity given the query vector, is a fundamental task in information retrieval, such as image retrieval [13, 55], web search [51, 52, 53], and item recommendation [32, 5, 33, 35, 48]. With a significant number of candidate vectors and vector dimensions, the exact nearest neighbor search becomes intractable due to the substantial computation costs and high query latency [42], resulting in a growing interest in Approximate Nearest Neighbors Search (ANNS).

Among the various indexes, the lightweight compressed indexes, especially those based on quantization, are notable for their substantial advantages in terms of low storage costs and efficient parallel processing. As a result, these indexes draw the attention from both academics and industries and are present in well-known open repositories such as FAISS [26]. Existing learning methods for

quantizers differ in whether there are explicit labels and whether they are related to queries. Traditional clustering-based methods for learning codebooks, e.g., PQ [24], OPQ [18], and AQ [3], are categorized as un-supervised algorithms only connected with document embeddings without the query information. Another type of technique [20] involves training the codebooks by optimizing the reconstruction-based loss, which relies on correlations with queries and documents but lacks explicit natural supervision signals. The last category of algorithms employs the query-dependent ground-truth labels to improve the retrieval performance, such as the existing interactions between query and candidates in LightRec [34], the ground-truth nearest neighbors in BLISS [21]. These additionally introduced supervision signals help capture the relationships between items and queries, thus contributing to better quantizers. However, the interaction label is only available in a small set of datasets and the expense of obtaining the ground-truth labels is particularly high, making the scenarios without any label information a more common case.

In this paper, we creatively propose a K nowledge D istillation for high dimensional search index framework ( KDindex for short) to distill knowledge from a more advanced teacher index model into lightweight indexes, under the circumstances without available label information. The advanced indexes, such as graph-based indexes, can yield more accurate results owing to their powerful expressiveness of the vector space and hence becomes able to teach the less-accurate lightweight indexes. Specifically, the top-k nearest results obtained from the teacher indexes act as the supervision signals to optimize the compressed functions. Therefore, index distillation is formulated as the top-k learning for the student index model from the teacher index. Specific ranking-oriented losses are exploited to learn the knowledge from the teacher model, which guides the results from the student model to have the same ranking orders as the teacher models. To avoid the trivial solutions, where the candidate vectors collapse to the identical centroid, we have two learning constraints for the learning objective. The reconstruction loss minimizes the distance between the original input and the compressed input, which pushes the centroids by the different query and candidate vectors. The second strategy is the posting list balance, which encourages an equal distribution of documents among each centroid. Thus, the candidates would be equally assigned to different centroids to avoid trivial solutions. Lastly, unlike recent iterative approaches that compulsively modify indexes, our approach uses a differentiable training process that updates the centroids and indexes simultaneously per mini-batch. This aims to minimize the error resulting from asynchronous updates. In this way, the student model would yield better performance thanks to the powerful expressiveness of the neighbor relationships of the teacher model, while keeping the low inference time and storage cost.

The contribution of this work can be summarized as follows:

To the best of our knowledge, this is the first attempt to extract knowledge from the high-precision search index with advanced structures into lightweight indexes in order to improve the retrieval performance of the lightweight index in high dimensions.

We propose the index distillation paradigm to learn the top-k nearest neighbors retrieved from the teacher models, where the top-k retrieved results act as the supervision signals to guide the learning of centroids, which is label-free and takes the query information into consideration.

We utilize ranking-oriented objectives as distillation loss with two learning constraints to avoid trivial solutions. Furthermore, we design the differentiable training process to avoid the asynchronous update of centroids and index assignments.

Experimental results on four benchmark datasets demonstrate that KDindex achieves a 40x index compression ratio, and 2x CPU speedup while maintaining comparable retrieval performance as the state-of-the-art compressed ANNS models.

2 Preliminaries

2.1 Problem Definition

Assume that there are M candidates { i j } j M =1 , such as documents, images, and items, in the retrieval system with no additional interaction information and the i -th candidate is represented by a D -dimensional vector d i ∈ R D . D usually appears to be huge numbers (larger than 100) in nowadays retrieval systems. Given a query vector q ∈ R D , the system attempts to return the top-k relevant candidates depending on the similarity scores S ( · , · ) between the query and candidate vectors, such as the inner product, L2 distance, and cosine distance.

Query

𝒒

Index

Top-K Ranked

documents list

𝒟𝐾

𝑇

= {𝒅𝑟1

, 𝒅

𝑟 2

, … , 𝒅𝐾}

Quantized

𝑄 𝒅𝒓𝟏

Quantized

𝑄 𝒅𝒓𝟐

Quantized

𝑄 𝒅𝑲

⋯

c. Student  Index

Codeword 1

Codeword 2

Codeword 3

Codeword w

Codewords

Posting List

Balanced

sinkhorn

-knopp

Step 1:

Derive index assignment of query

and documents via quantizer.

Step 2:

Update codewords and posting list by

Reconstruction loss and Distillation loss.

top-K ranked candidates

Loss

2.2 Lightweight Compressed Indexes

With the increasing number of candidates, Approximate Nearest Neighbors Search (ANNS) and Maximum Inner Product Search (MIPS) strategies are widely used to accelerate the search speed at the expense of accuracy, among which lightweight compressed indexes (hashing and quantization) are widely used because of their strong performance in terms of both effectiveness and efficiency.

The quantization-based indexes, a representative compressed ANNS algorithm, quantize each candidate vector into the reconstruction vector from the B codebooks, each of which contains W vectors (codewords). We denote the b -th codebook as C b and the w -th codeword in the b -th codebook as c b w . The Additive Quantization (AQ) encodes the vector as a sum of B codewords, i.e., ˜ d = Q ( d ) = ∑ B b =1 c b w b , while the Product Quantization (PQ) and Optimized product quantization (OPQ) encode the vector as the concatenation of B codewords of each subspace, i.e., ˜ d = Q ( d ) = c 1 w 1 ⊕ c 2 w 2 ⊕ ... ⊕ c B w B . The codeword with the minimal distance from the original vector is chosen, i.e., c b w = arg min i ∈{ 1 , 2 ,...,W } ∥ x -c b w i ∥ 2 . Q ( · ) denotes the quantizer function. The codewords are usually learned by the unsupervised clustering algorithms or the reconstruction loss, and the posting lists which record the belonging candidates are then constructed. During the inference stage, the similarity scores between the query vector and the codewords are calculated and the approximate nearest candidates are oriented. Consequently, given the well-trained indexes, the overall inference time complexity is only correlated with the number of codewords and codebooks.

Learning to hash is also an effective method to compress high-dimensional data into low-dimensional binarized codes, similar to those clustering-based conventional quantization methods. Learning to hash approaches have the property that objects that are close to each other have a higher probability of colliding than objects that are far apart across various distance metrics. The drawbacks of these approaches are the requirement for a large number of hash tables in order to achieve good search quality and these methods are unmindful of the distribution of vectors, often leading to lop-sided partitions and long query times. Thus, we focus more on the quantization-based compressed methods.

3 Index Distillation

Despite the efficiency advantages of lightweight quantization-based indexes, they are still limited by the representation space of codewords. Existing works introduce additional ground-truth labels to supervise the learning. For example, LightRec [34] exploits the interaction information between the certain query and document, and BLISS [21] takes the ground-truth neighbors over candidates to enhance the discrimination probability. These extra signals benefit better capturing the querydependent similarity relationships between candidates. However, in more circumstances, there are no ground-truth labels considering the high expense to obtain them. Thus, we focus on scenarios without label information and aim to enhance the performance of lightweight indexes.

Several existing works introduce the knowledge of query to adjust the choice of the centroids [20] or to randomly sample neighbors to learn the centroids. These works provide less relevant candidates, thus resulting in less accurate inference results. Intuitively, if we use more progressive index structures, such as graph-based methods, to encode the similarity information over the whole vector space, more accurate results would be obtained. This motivates the progressive indexes as the teacher to distill more informative high-dimensional knowledge to the lightweight quantization-based indexes, with the goal of improving the accuracy of lightweight indexes under high dimensions.

3.1 Overview of Index Distillation

The whole framework, i.e., Knowledge Distillation for high dimensional search index (KDindex), consists of a high-precision teacher search model, which is well-learned, and a lightweight student search model to be learned, where the student model attempts to learn knowledge from the top-k relevant candidates retrieved from the teacher model. Specifically, given the well-trained candidate vectors, when a query vector q requests, the teacher search model returns the set of the top-k approximate nearest neighbors D T K = { i | rank ( i | q ) ≤ K } . These candidates then behave as weak signals to supervise the learning of codewords in the student models, where the student indexes are guided to learn the same ranking orders and finally to return the similar indexes for the given query and candidate. Figure 1 visualizes the whole training process of KDindex, with the three key components: Distillation Loss, two Learning Constraints, and a Differentiable Training paradigm which are detailed in the following sections. The whole training process is illustrated in Algorithm 2.

3.2 Initialization

The initialization of the student search model involves the following steps: (1) Initialize the centroids (codewords) and assign indexes (2) Update the centroids based on the distillation loss and reconstruction loss given the training queries (3) Adjust the posting lists depending on the latest centroids. A straightforward process is to update the posting lists per T epochs after centroids updating is finished by enumerating all training queries, similar to the paradigm of BLISS [21].

3.3 Distillation Loss

Knowledge distillation was first proposed for classification tasks [22], where the probabilities of each class attained from the large-scale teacher network are considered as soft labels to supervise the learning of the small-size student network. The cross-entropy loss is commonly used as the distillation loss to minimize the difference between the teacher and student networks. Here, the teacher search model provides the top-k relevant candidates rather than the continuous value of probabilities. We exploit the following Distributed-based loss as the distillation loss to guide the student indexes to return the same nearest results. This objective function follows as:

where B denotes the number of codebooks and W is the number of codewords in each codebook. w i = 1 rank ( i ) corresponds to the top-k list given from the teacher model. p q bk denotes the similarity score between the query q and the codeword c b k , i.e., p q bk = S ( q , c b k ) , and p d i bk denotes the similarity score between the candidate d i and the codeword c b k , i.e., p d i bk = S ( d i , c b k ) . The normalized value ˜ p q bk and ˜ p d i bk are calculated over the W codewords for each codebook through the softmax function.

This loss attempts to minimize the distance between the queries and top-k neighbors by calculating the similarity scores with all the centroids. Thus, we could obtain more information from centroids and focus on the top-K nearest neighbors.

3.4 Learning Constraints

Although the distillation loss supervises the learning of quantizers, the student model would collapse with a trivial solution where all candidates are assigned to the identical centroid. To avoid trivial solutions, we propose two strategies and we will introduce them one by one.

Reconstruction Loss. A straightforward solution is to design the reconstruction loss to draw the distance between the input vector and the reconstructed vector encoded through the quantized function. Specifically, MSE loss is adopted here to specialize the reconstruction loss for the both query and candidate vectors:

where x ∈ R D represents the original input vector such as the query vector q and the candidate vector d . Q denotes the quantized function to reconstruct the input vector and the corresponding codewords would be updated through the loss.

Balanced Posting List. The posting list refers to the candidate list according to the certain centroids. The trivial solutions would result in an imbalanced posting list, where some posting lists would be substantially longer than others and thus resulting linear complexity in the worst case. Thus, we impose a balanced clustering constraint to guide the candidate vectors to be equally assigned to all quantization centroid embeddings. Here, we take one codebook as an example to formulate the learning objective with the constraint:

where q ( k | d i ) denotes the binary value indicating whether the candidate vector is quantized to the k -th centroid and the second condition ensures that only one centroid is selected. The last condition guides each centroid to be allocated an equal number of candidates to achieve a balanced posting list. Thus, the indexes for different candidates would be varied to avoid the trivial solution.

To solve Eq (3), we relax the binary constraint q ( k | d i ) to the continuous probability, i.e., q ( k | d i ) ∈ (0 , 1) , to get the approximated solution. The problem then can be reduced to the optimal transport problem by taking the distance as the cost of quantization and we use the Sinkhorn-Knopp algorithm [10] to solve it. Algorithm 1 details the balance strategy to uniformly allocate the candidates.

3.5 Differentiable Training

The re-assignment of the indexes, including ' arg min ' operation, avoids the consecutive computation and thus interrupts the backward of the gradient, resulting in the iterative training with the two separate steps. However, the update of the index is later than the update of embedding, resulting in a large deviation in the quantized vector and inaccuracy of the retrieval performance. To keep the index assignment up-to-date during each training batch, we attempt to design a differentiable training manner for continuous computation. The Gumbel-Softmax Trick [23] is a common and popular solution, motivated by which we adopt a similar sample operation to copy the gradient into the one-hot embeddings. Surprisingly, the probabilities of each centroid are calculated in the

aforementioned balanced posting list and can be regarded as the weight of the centroids for updating, as shown in Line 9-11 of Algorithm 2.

The whole training process is illustrated in Algorithm 2, which can be computed in parallel for each mini-batch. Accordingly, the learning of centroids and the index re-assignment are updated within the same mini-batch, rather than the asynchronous epochs.

3.6 Complexity Analysis

The training ways of KDindex and basic Qutization methods (AQ, PQ and OPQ) are different, therefore the complexity is different in the training phase. The indexing and inference way of KDindex and Qutization methods are consistent and there is no additional time complexity in the index and inference stages. More details are described in Table 1.

Table 1: The time complexity of KDindex, HNSW, and ScaNN. Denoted by D item embedding dimension, B the number of subspaces, W the number of centroids in each codebook, M the batch size (the number of queries in each batch), and K the number of neighbors. N is the number of items. As for ScaNN, K v denotes the number of centroids in VQ ( vector quantization ) and K p in PQ ( Product Quantization ).

4 Experiment

4.1 Experiment Setup

Datasets and Metrics. Four large-scale retrieval benchmarks, including SIFT1M , GIST1M from ANN datasets [2], MS MARCO Doc and MS MARCO Passage from the TREC 2019 Deep Learning

Track [9], are used to validate the effectiveness of the proposed KDindex. SIFT1M and GIST1M both have exactly 1 Million database points, and the dimension is 128 and 960, respectively. Document Retrieval consists of 3.2 M documents, 0.36 M training queries, and 5 K development queries. Passage Retrieval has a corpus of 8.8 M passages, 0.8 M training queries, and 0.1 M development queries. All the datasets offer the set of training queries and testing queries. The detailed specifications of the datasets are shown in Table 6. The search model returns the topK retrieval results given the test query and thus we use ranking-based metrics to evaluate the performance, including Recall, NDCG and MRR with a cutoff of 10.

Implementation Details. SIFT1M and GIST1M datasets contain the candidate vectors with the dimension of 128 and 960 respectively. The vectors for MS MARCO Doc and MS MARCO Passage are generated based on huggingface transformers [47] and the dimension is set to 768 following previous works [14]. The similarity function S for SIFT1M and GIST1M is the Euclidean distance and for MS MARCO Doc and Passage is the inner product. We exploit a well-trained HNSW [37] model as the teacher model and K = 10 nearest neighbors are retrieved accordingly. We adopt PQ, OPQ, and AQ as the student model for KDindex. For a fair comparison, we run current supervised methods with optimized objectives proposed in KDindex instead of interaction information-based loss function. Each vector is quantized by B = 8 codebooks, each of which contains W = 256 codewords by default. The centroids are trained with a learning rate of 0.01 and optimized by the Adam [28] optimizer. The batch size is set to 64.

4.2 Overall performance

We first compare the retrieval performance within the same retrieval time to verify the effectiveness of KDindex. Experiments are run 5 times to conduct the t-test and the performance of the retrieval is illustrated in Table 2, from which we have the following findings:

KDindex respectively improves the retrieval performance of three native student models, demonstrating the general effectiveness of KDindex. KDindex has a relative 10.49% and 9.65% improvement in

32.5

34.75

37.0

Recall@10

34.50

36.5

34.25

36.0

31.5

34.00

35.5

0

20

40

60

80

100

0

20

40

60

80

100

120

0

20

40

60

80

100

Epoch

Epoch

Epoch

Epoch-Recall MS MARCO Doc PQ

Epoch-Recall MS MARCO Doc OPQ

Epoch-Recall MS MARCO Doc AQ

19.0

F8

16.5

18.5

Recall@10

0

16.0

10

6

18.0

15.5

17.5

★

KDindex_differentiable

4

15.0

17.0

★

KDindex_iterative

Quantization

0

25

50

75

100

125

150

175

0

20

40

60

80

100

120

0

20

40

60

80

100

Epoch

Epoch

Epoch

Recall@10 over all datasets compared with the native OPQ and AQ, respectively. This demonstrates that KDindex receives the supervision signals of the top-k retrieved results and thus more neighbour relationships are distilled for the student, resulting in better performance of item recall.

KDindex shows better performance than learnable quantization-based indexes with no ground-truth labels. By choosing appropriate quantization-based indexes, KDindex would outperform all baseline methods, especially the learnable indexes DeepPQ, demonstrating the advantage of the proposed KDindex. The introduction of distillation loss has an edge over capturing the relationships between candidates, where the top-k retrieved candidates provide stronger signals for learning.

The improvement of KDindex is more significant when the distribution between query space and database space is different. KDindex obtains more improvements on MS MARCO Doc and MS MARCO Passage datasets which search documents by the inner product. The similarity could not be obtained by original quantization methods. Experiments have proved that KDindex could indeed learn the neighbor relationship between query vectors and candidates by the distillation loss function.

4.3 Performance of Differentiable Training

It is extremely difficult for the model to learn the codebooks as well as the index at the same time during the initialization phase in a differentiable training manner. Thus, we perform experiments by warming up the codebooks by Initialization and we get the following results in terms of Recall@10 on four datasets as Fig. 2. We adopt the early stop strategy to get the best model.

Initialization. We obtain the pre-trained codebooks by iterative training manners and continue differentiable training when the index assignment is approaching being balanced ( max |P i | -min |P j | < N W , i, j ∈ W ) where |P i | denotes the length of the i -th posting list. To accelerate the iterative training, codebooks are warmed by original quantization methods such as PQ, OPQ, and AQ.

Findings. KDindex converges to a better solution through the differential training manner. Within the dozens of epochs, the index assignment of iterative training becomes balanced, which warms up the centroids for later easier learning and thus relatively reduces the learning difficulty for both codebooks and indexes. Starting from this point, KDindex with differentiable training consecutively outperforms that with iterative training, which achieves a relative improvement of 1.63% in terms of Recall@10 on both datasets, demonstrating the effectiveness of synchronizing updates for codebooks and indexes. As for the different quantization functions, the improvements of Recall@10 among different student models (PQ, OPQ, and AQ) are 1.49%, 1.46%, and 1.94%, respectively. The better performance of KDindex(AQ) may be attributed to its better expressiveness with more parameters. Finally, the improvement of Recall@10 on MS MARCO Doc by KDindex(PQ) is 0.40%, which is smaller than the other model since the express ability of PQ is limited. The improvement of Recall@10 on SIFT1M by KDindex(AQ) is 0.39% since the express ability of AQ is strong on the L2 distance dataset and no more improvements can be obtained easily.

4.4 Memory Efficiency

As shown in Table 3, KDindex is more memory-efficient than the non-exhaustive ANNS method, HNSW, while performing comparable effectiveness. The index size of KDindex is less than 0.6GB, while HNSW requires storing all the original vectors and the graphs and consumes 0.5GB, 3.8GB, 9.2GB,

and 25.3GB on four datasets respectively for storing indexes. The average Compression Ratio between the memory of the non-exhaustive method (HNSW) and KDindex on four datasets approximates 40x, demonstrating the superiority of the low memory cost of KDindex owing to the lightweight structure.

4.5 Hyperparameter Analysis

In this section, we analyze the effects of hyperparameters, including the number of nearest neighbors retrieved from the teacher model K , the number of codebooks B , and the number of codewords in each codebook W .

Effects of B and W Figure 3 shows the effects of B and W , where B varies in { 4 , 8 , 16 , 32 } and W in { 64 , 128 , 256 , 512 } . We conduct experiments on SIFT1M and compare the results of KDindex with different quantization methods.

Firstly, we see that, in terms of task performance, decreasing W has a far more traumatizing impact on KDindex. This is due to the nearest neighbor estimate losing accuracy when the number of codewords in each codebook drops.

Secondly, increasing B or W would typically improve the task performance at the expense of lower efficiency and more storage cost, which means one can adjust B and W in the meantime to obtain a trade-off between effectiveness and efficiency.

Thirdly, we notice that the small W and large B combination is preferable

PQ

OPQ

AQ

-14.62

18.29

26.13

31.80

14.65

19.07

26.45

32.33

14.89

20.73

28.87

34.96

128

18.14

30.95

37.65

37.94

18.27

33.48

39.45

39.82

19.88

35.19

41.53

41.76

256

19.74

31.27

42.01

42.10

20.57

33.85

44.01

44.60

21.65

35.48

46.14

46.23

512

20.69

32.48

44.10

44.35

21.64

34.05

46.51

46.73

22.48

35.72

48.49

48.60

4

8

16

32

4

8

T

16

32

4

8

16

32

B

B

B

-14.74

18.63

26.40

32.46

14.79

20.24

27.20

34.23

14.61

21.71

30.26

36.72

128

18.23

32.18

37.92

39.10

19.48

34.39

40.33

40.91

20.90

36.90

43.48

43.88

W

256

19.94

32.53

42.56

43.30

21.39

34.77

45.21

45.21

22.87

37.30

48.50

48.50

512

20.81

32.73

44.79

45.57

22.53

34.89

47.49

47.53

23.85

37.53

50.9351.10

4

8

16

32

4

8

16

32

4

8

16

32

B

B

B

to the opposite way around. For instance, among the three quantization methods with KDindex, ( B = 32; W = 64) performs better than ( B = 4; W = 512) on the SIFT1M.

Effects of K The relevant results come from the teacher model directly influence the distribution between the teacher and student and thus play a vital role in distillation loss. We vary the number of the nearest neighbors among { 5 , 10 , 20 , 50 , 100 } and report the results in Table 4.

The model performs best when the number of neighbors is equal to 10, which is consistent with the cutoff. When the number of neighbors increases, leading to more data points are trained, the effects

degrade instead, which is contrary to the intuition that more candidates get better. The reason may lie in that the lower-ranked items have lower relevance with lower confidence and may bring more noise in training compared with the high-ranked items.

4.6 Ablation Study

KDindex includes three strategies, i.e., distillation loss, balance strategy, and differentiable training. We conduct an ablation study by incrementally adding the strategies to the basic quantization methods (PQ, OPQ, and AQ) to investigate their contributions. Specifically, we use the following model variants: (1) Quantization [24, 18, 3]: They are general quantization methods that serve as the backbone of the student model in KDindex. (2) Initialization is warmed up by quantization methods. It updates index assignments and centroids iteratively. is warmed up by quantization methods. It updates index assignments and centroids iteratively. (Details can be found in Section 4.4. We obtain the pre-trained codebooks by iterative training manners. ) (3) w/o Distillation loss denotes the training without knowledge distilled from the teacher model (HNSW in the experiment). It optimizes the centroids and trains the encoder under the constraint of reconstruction loss and balance strategy. (4) w/o Balance strategy denotes methods that without Sinkhorn-Knopp balance strategy. The reconstruction loss is so necessary to avoid a trivial solution that could not be removed. (5) KDindex denotes methods that differentially train models with Reconstruction loss, Distillation loss, and Balance strategy.

We conduct experiments on the MS MARCO Doc dataset with the same settings as mentioned above and report MRR@10 and Recall@10. Results are shown in Tab. 5. We can see that all three strategies contribute to the effectiveness of KDindex. With the help of the distillation loss, Recall@10 and MRR@10 achieve 14.34% and 10.01% relative improvements respectively. This indicates that the knowledge distilled based on the ranking list retrieved from the teacher model contributes to the learning of the centroids in the student model. Furthermore, the posting list balance strategy leads to a relative 2.61% and 2.76% improvement on Recall@10 and MRR@10, demonstrating that the avoidance of trivial solutions by controlling the length of the posting list is beneficial to building more appropriate centroids.

5 Conclusion

This paper presents KDindex, a novel framework for learning compressed indexes in high dimensions, which distills knowledge of the top-k relevant documents from the teacher index into the lightweight indexes to improve the retrieval accuracy for the student search model. To guide the student model to learn the same ranking list, rank-oriented distillation losses are designed and the posting list balance strategy is proposed to avoid the trivial solutions implemented with the Sinkhorn-Knopp algorithm. The learning of the student search model consists of the centroids embedding learning and the index assignment is exploited to realize the end-to-end differentiable learning. We conduct experiments on four publicly available benchmarks, where KDindex achieves impressive performance. Even if KDindex utilizes the approximated nearest neighbors and compresses the index by 40x, it outperforms some competitive models regarding ground truth generated by bruce-force as labels and gains better ranking performance than the state-of-the-art one.

Acknowledgments and Disclosure of Funding

The work was supported by grants from the National Natural Science Foundation of China (No.62022077 and 61976198).

References

Martin Aumüller, Erik Bernhardsson, and Alexander Faithfull. Ann-benchmarks: A benchmarking tool for approximate nearest neighbor algorithms. In Similarity Search and Applications: 10th International Conference, SISAP 2017, Munich, Germany, October 4-6, 2017, Proceedings 10 , pages 34-49. Springer, 2017.

Martin Aumüller, Erik Bernhardsson, and Alexander Faithfull. Ann-benchmarks: A benchmarking tool for approximate nearest neighbor algorithms. Information Systems , 87:101374, 2020.

Artem Babenko and Victor Lempitsky. Additive quantization for extreme vector compression. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition , pages 931-938, 2014.

Christopher Burges, Robert Ragno, and Quoc Le. Learning to rank with nonsmooth cost functions. Advances in neural information processing systems , 19, 2006.

Jin Chen, Defu Lian, Binbin Jin, Xu Huang, Kai Zheng, and Enhong Chen. Fast variational autoencoder with inverted multi-index for collaborative filtering. In Proceedings of the ACM Web Conference 2022 , pages 1944-1954, 2022.

Ting Chen, Lala Li, and Yizhou Sun. Differentiable product quantization for end-to-end embedding compression. In International Conference on Machine Learning , pages 1617-1626. PMLR, 2020.

Ting Chen, Martin Renqiang Min, and Yizhou Sun. Learning k-way d-dimensional discrete codes for compact embedding representations. In International Conference on Machine Learning , pages 854-863. PMLR, 2018.

Yuntao Chen, Naiyan Wang, and Zhaoxiang Zhang. Darkrank: Accelerating deep metric learning via cross sample similarities transfer. In Proceedings of the AAAI conference on artificial intelligence , volume 32, 2018.

Nick Craswell, Bhaskar Mitra, Emine Yilmaz, Daniel Campos, and Ellen M Voorhees. Overview of the trec 2019 deep learning track. arXiv preprint arXiv:2003.07820 , 2020.

Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems , 26, 2013.

Xinyan Dai, Xiao Yan, Kelvin KW Ng, Jiu Liu, and James Cheng. Norm-explicit quantization: Improving vector quantization for maximum inner product search. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 34, pages 51-58, 2020.

Mayur Datar, Nicole Immorlica, Piotr Indyk, and Vahab S Mirrokni. Locality-sensitive hashing scheme based on p-stable distributions. In Proceedings of the twentieth annual symposium on Computational geometry , pages 253-262, 2004.

Ritendra Datta, Dhiraj Joshi, Jia Li, and James Z Wang. Image retrieval: Ideas, influences, and trends of the new age. ACM Computing Surveys (Csur) , 40(2):1-60, 2008.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 , 2018.

Chao Feng, Wuchao Li, Defu Lian, Zheng Liu, and Enhong Chen. Recommender forest for efficient retrieval. Advances in Neural Information Processing Systems , 35:38912-38924, 2022.

Chao Feng, Defu Lian, Xiting Wang, Zheng Liu, Xing Xie, and Enhong Chen. Reinforcement routing on proximity graph for efficient recommendation. ACM Transactions on Information Systems , 41(1):1-27, 2023.

Lianli Gao, Xiaosu Zhu, Jingkuan Song, Zhou Zhao, and Heng Tao Shen. Beyond product quantization: Deep progressive quantization for image retrieval. arXiv preprint arXiv:1906.06698 , 2019.

Tiezheng Ge, Kaiming He, Qifa Ke, and Jian Sun. Optimized product quantization. IEEE transactions on pattern analysis and machine intelligence , 36(4):744-755, 2013.

Ruiqi Guo, Sanjiv Kumar, Krzysztof Choromanski, and David Simcha. Quantization based fast inner product search. In Artificial intelligence and statistics , pages 482-490. PMLR, 2016.

Ruiqi Guo, Philip Sun, Erik Lindgren, Quan Geng, David Simcha, Felix Chern, and Sanjiv Kumar. Accelerating large-scale inference with anisotropic vector quantization. In International Conference on Machine Learning , pages 3887-3896. PMLR, 2020.

Gaurav Gupta, Tharun Medini, Anshumali Shrivastava, and Alexander J Smola. Bliss: A billion scale index using iterative re-partitioning. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining , pages 486-495, 2022.

Geoffrey Hinton, Oriol Vinyals, Jeff Dean, et al. Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531 , 2(7), 2015.

Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. arXiv preprint arXiv:1611.01144 , 2016.

Herve Jegou, Matthijs Douze, and Cordelia Schmid. Product quantization for nearest neighbor search. IEEE transactions on pattern analysis and machine intelligence , 33(1):117-128, 2010.

Yunjiang Jiang, Han Zhang, Yiming Qiu, Yun Xiao, Bo Long, and Wen-Yun Yang. Givens coordinate descent methods for rotation matrix learning in trainable embedding indexes. arXiv preprint arXiv:2203.05082 , 2022.

Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion-scale similarity search with GPUs. IEEE Transactions on Big Data , 7(3):535-547, 2019.

SeongKu Kang, Junyoung Hwang, Wonbin Kweon, and Hwanjo Yu. De-rrd: A knowledge distillation framework for recommender system. In Proceedings of the 29th ACM International Conference on Information & Knowledge Management , pages 605-614, 2020.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Jon Kleinberg. The small-world phenomenon: An algorithmic perspective. In Proceedings of the thirty-second annual ACM symposium on Theory of computing , pages 163-170, 2000.

Jae-woong Lee, Minjin Choi, Jongwuk Lee, and Hyunjung Shim. Collaborative distillation for top-n recommendation. In 2019 IEEE International Conference on Data Mining (ICDM) , pages 369-378. IEEE, 2019.

Wen Li, Ying Zhang, Yifang Sun, Wei Wang, Mingjie Li, Wenjie Zhang, and Xuemin Lin. Approximate nearest neighbor search on high dimensional data-experiments, analyses, and improvement. IEEE Transactions on Knowledge and Data Engineering , 32(8):1475-1488, 2019.

Defu Lian, Yong Ge, Fuzheng Zhang, Nicholas Jing Yuan, Xing Xie, Tao Zhou, and Yong Rui. Content-aware collaborative filtering for location recommendation based on human mobility data. In 2015 IEEE international conference on data mining , pages 261-270. IEEE, 2015.

Defu Lian, Qi Liu, and Enhong Chen. Personalized ranking with importance sampling. In Proceedings of The Web Conference 2020 , pages 1093-1103, 2020.

Defu Lian, Haoyu Wang, Zheng Liu, Jianxun Lian, Enhong Chen, and Xing Xie. Lightrec: A memory and search-efficient recommender system. In Proceedings of The Web Conference 2020 , pages 695-705, 2020.

Defu Lian, Yongji Wu, Yong Ge, Xing Xie, and Enhong Chen. Geography-aware sequential location recommendation. In Proceedings of the 26th ACM SIGKDD international conference on knowledge discovery & data mining , pages 2009-2019, 2020.

Jie Liu, Xiao Yan, Xinyan Dai, Zhirong Li, James Cheng, and Ming-Chang Yang. Understanding and improving proximity graph based maximum inner product search. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 34, pages 139-146, 2020.

Yu A Malkov and Dmitry A Yashunin. Efficient and robust approximate nearest neighbor search using hierarchical navigable small world graphs. IEEE transactions on pattern analysis and machine intelligence , 42(4):824-836, 2018.

Stanislav Morozov and Artem Babenko. Non-metric similarity graphs for maximum inner product search. Advances in Neural Information Processing Systems , 31, 2018.

Marius Muja and David G Lowe. Scalable nearest neighbor algorithms for high dimensional data. IEEE transactions on pattern analysis and machine intelligence , 36(11):2227-2240, 2014.

Behnam Neyshabur and Nathan Srebro. On symmetric and asymmetric lshs for inner product search. In International Conference on Machine Learning , pages 1926-1934. PMLR, 2015.

Yingqi Qu, Yuchen Ding, Jing Liu, Kai Liu, Ruiyang Ren, Wayne Xin Zhao, Daxiang Dong, Hua Wu, and Haifeng Wang. Rocketqa: An optimized training approach to dense passage retrieval for open-domain question answering. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies , pages 5835-5847, 2021.

Hanan Samet. Foundations of multidimensional and metric data structures . Morgan Kaufmann, 2006.

Anshumali Shrivastava and Ping Li. Asymmetric lsh (alsh) for sublinear time maximum inner product search (mips). Advances in neural information processing systems , 27, 2014.

Jiaxi Tang and Ke Wang. Ranking distillation: Learning compact ranking models with high performance for recommender system. In Proceedings of the 24th ACM SIGKDD international conference on knowledge discovery & data mining , pages 2289-2298, 2018.

Haoyu Wang, Defu Lian, and Yong Ge. Binarized collaborative filtering with distilling graph convolutional networks. arXiv preprint arXiv:1906.01829 , 2019.

Jun Wang, Wei Liu, Sanjiv Kumar, and Shih-Fu Chang. Learning to hash for indexing big data-a survey. Proceedings of the IEEE , 104(1):34-57, 2015.

Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-theart natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations , pages 38-45, Online, October 2020. Association for Computational Linguistics.

Chenwang Wu, Defu Lian, Yong Ge, Zhihao Zhu, and Enhong Chen. Influence-driven data poisoning for robust recommender systems. IEEE Transactions on Pattern Analysis and Machine Intelligence , 2023.

Hanwei Wu and Markus Flierl. Learning product codebooks using vector-quantized autoencoders for image retrieval. In 2019 IEEE Global Conference on Signal and Information Processing (GlobalSIP) , pages 1-5. IEEE, 2019.

Xiang Wu, Ruiqi Guo, Ananda Theertha Suresh, Sanjiv Kumar, Daniel N Holtmann-Rice, David Simcha, and Felix Yu. Multiscale quantization for fast similarity search. Advances in neural information processing systems , 30, 2017.

Shitao Xiao, Zheng Liu, Weihao Han, Jianjin Zhang, Defu Lian, Yeyun Gong, Qi Chen, Fan Yang, Hao Sun, Yingxia Shao, et al. Distill-vq: Learning retrieval oriented vector quantization by distilling knowledge from dense embeddings. In Proceedings of the 45th International ACM SIGIR Conference on Research and Development in Information Retrieval , pages 1513-1523, 2022.

Jingtao Zhan, Jiaxin Mao, Yiqun Liu, Jiafeng Guo, Min Zhang, and Shaoping Ma. Jointly optimizing query encoder and product quantization to improve retrieval performance. In Proceedings of the 30th ACM International Conference on Information & Knowledge Management , pages 2487-2496, 2021.

Jingtao Zhan, Jiaxin Mao, Yiqun Liu, Jiafeng Guo, Min Zhang, and Shaoping Ma. Learning discrete representations via constrained clustering for effective and efficient dense retrieval. In Proceedings of the Fifteenth ACM International Conference on Web Search and Data Mining , pages 1328-1336, 2022.

Ting Zhang, Chao Du, and Jingdong Wang. Composite quantization for approximate nearest neighbor search. In International Conference on Machine Learning , pages 838-846. PMLR, 2014.

Zaixi Zhang and Qi Liu. Learning subpocket prototypes for generalizable structure-based drug design. In Proceedings of the 40th International Conference on Machine Learning , ICML'23, 2023.

A Related Work

A.1 ANNs and MIPS

ANNS achieves highly efficient vector search by allowing a small number of errors. Generally, there are two kinds of ANNS algorithms: non-exhaustive ANNS methods [39, 36, 38, 29] and vector compression methods [12, 3, 20]. Specifically, Non-exhaustive ANNS methods do not compress the index. They reduce the number of candidates for each query to speed up the retrieval process. Popular algorithms include tree search [16, 15] and graph search [36, 38, 29]. Vector compression methods mainly aim to compress the index to accelerate retrieval. Popular algorithms include hashing [12, 40, 43] and quantization [3, 20, 19].

Under the constraints of storage, compressed methods are widely investigated by researchers. Product quantization [24, 7] decomposes the vector representation space into the Cartesian product of subspaces. Optimized product quantization (OPQ) [18] jointly learns space decomposition and subspace quantization. Multi-scale Quantization [50] includes a multi-scale framework so that it can learn a separate scalar quantizer. Composite Quantization [54] and Additive Quantization [3] do not decompose space but directly learn multiple codebooks. There are also some algorithms that take query information into account. NEQ [11] decomposes the quantization error into norm error and direction error and improves existing VQ techniques for MIPS. ScaNN [20] computes the weight for each pair of vectors. Different from NEQ and ScaNN, KDindex utilizes query and corresponding top-k candidates. LSH (Local sensitive hashing) is a data-independent un-supervised method, similar to those clustering-based conventional quantization methods. BLISS [21] regards ground truth as labels. However, the ground truth is difficult to obtain in huge quantities of databases. LSH approaches have the property that objects that are close to each other have a higher probability of colliding than objects that are far apart across various distance metrics. The drawbacks of these approaches are the requirement for a large number of hash tables in order to achieve good search quality and these methods are unmindful of the distribution of vectors, often leading to lop-sided partitions and long query times. Interested readers could refer to the surveys [46, 31].

A.2 Knowledge Distillation

Knowledge Distilling (KD) was first proposed in [22], in which a complex neural network was firstly trained and then transferred to a small model. Following this, DarkRank [8] proposed a method combining deep metric learning and Learning to rank technique with KD to solve image retrieval and image clustering tasks. In addition, a few recent methods [30, 41] have adopted knowledge distillation to RS. RD [44] firstly proposes a KD method that makes the student give high scores on the top-ranked items of the teacher's recommendation list. Similarly, CD [30] makes the student imitate the teacher's prediction scores with particular emphasis on the items ranked highly by the teacher. The most recent work, RRD [27], formulates the distillation process as a relaxed ranking matching problem between the ranking list of the teacher and that of the student. However, there are limited works focusing on index building under knowledge distillation.

There are also numerous works that use knowledge distillation to improve the performance of hashingbased [45] and quantization-based [51] codes, The ranking information is distilled from a graph-based network to enhance the performance of hashing codes. However, these works rely on the ground-truth labels (user-item interactions) to learn the ranking orders. This is different from our work, where label information is not accessible for learning. The quantization based method uses knowledge distillation for ranking candidates in web search tasks and applies the sampling technique to rank a sample of the document from data each time. But this technique is not applicable to training a ranking model when documents and queries are represented with no content information. In this case, the labeled model training cannot be easily generalized to all documents and queries. In contrast, KDindex relaxes the requirement on labeled data and can be trained purely with unlabeled data.

B More Details of Experimental Settings

B.1 Dataset Statistics

B.2 Baselines

The two groups of baseline ANNS models are compared to KDindex.

The first group is Non-quantization-based ANNS methods , which accelerate the search by reducing the number of candidates. BLISS [21] adopts the learning-to-index framework to learn the hashing-based compressed functions. ScaNN [20] quantizes with anisotropic quantization loss functions which greatly penalizes the parallel component of a datapoint's residual relative to its orthogonal component. HNSW [37] builds a hierarchical set of proximity graphs. Results can be found in Appendix D.

The second is Quantization-based ANNS methods , which compress the embeddings by hashing or quantization functions. PQ [24] decomposes the vector representation space into the Cartesian product of subspaces. OPQ [18] jointly learns space decomposition and subspace quantization. AQ [3] represents each vector as a sum of several components each coming from a separate codebook. The baselines are implemented based on the Faiss ANNS library [26]. The parameters B and W are set to be the same as KDindex. DiffPQ [6], differentiable product quantization, a generic and end-to-end learnable compression framework. DeepPQ [17], deep progressive quantization, end-toend learns the quantization codes sequentially. PQ-VAE [49], an unsupervised model for learning discrete representations by combining product quantization and auto-encoders. The CNN blocks are replaced with MLP because the image datasets have been extracted as 512-dimension features. GCD [25] learns rotation matrix via a family of block Givens coordinate descent algorithms. RepCONC [53] requires data points to be uniformly clustered around the quantization centroids.

B.3 Implemental details

Teacher model is instantiated as HNSW. The details are described as Tab. 7, where M denotes the number of neighbors each node, efConstruction denotes expansion factor at construction time and efSearch denotes expansion factor at search time. To obtain good recall performance, M, efConstruction and efSearch are tuned. Building and Searching are on multi-threads for train sets.

C Varying Distillation Loss

C.1 Distillation Losses

Knowledge distillation was first proposed for classification tasks, where the probabilities of each class attained from the large-scale teacher network are considered as soft labels to supervise the learning

of the small-size student network. The cross-entropy loss is commonly used as the distillation loss to minimize the difference between the teacher and student networks. Here, the teacher search model provides the top-k relevant candidates rather than the continuous value of probabilities. Thus, three ranking-oriented losses are designed to distill knowledge from the more accurate indexes to guide the student indexes to return the same nearest results.

Lambdarank loss : The pair-wise ranking-based loss is widely used to learn the ranking list by leading the high-ranked candidate to have higher similarity scores. Lambdarank [4] further introduces the change of the indicators, e.g., NDCG, to put more attention on more important candidates that have not been well ranked. The loss follows as:

where D T K denotes the top-k results retrieved from the teacher model and p i = S ( q , Q ( d i )) is the similarity score between the query vector and the quantized vector of the candidate i . Q is the quantizer function related to the codebooks C . ∆ NDCG @10 ij denotes the change with respect to NDCG@10 if changing the i -th ranked and j -th ranked candidate.

Weighted KL loss : Similar to the class distribution in classification tasks, the similarity distribution over the top-k retrieved candidates can also be obtained. One is based on the ground-truth vectors and the other one is based on the quantized vectors. In order to ensure the ranking orders correspond to the top-k list, the rank information is also considered where the high-ranked items are more concerned. Finally, the loss function follows as:

where ˜ p i denotes the normalized ranked similarity score with the quantized vector and ˜ p g i with the ground-truth vector. Specifically,

˜ p i and ˜ p g i are the normalized values over the top-k retrieved candidates depending on the softmax function. w i = 1 rank ( i ) denotes the ranking weight according to the ranking orders among the top-k results from the teacher model. The weighted KL loss attempts to minimize the distance between the ground-truth vector and the quantized vector for the top-k relevant candidates to learn better centroids. The introduced rank-oriented weight further guides the student index to return the same ranking list.

Distributed-based loss : Instead of being oriented by the score between query and candidates as above, we attempt to minimize the distance between the queries and top-k neighbors by calculating the similarity scores with all the centroids. Thus, we could obtain more information from centroids and focus on the top-K nearest neighbors. The distributed-based loss function follows as:

where B denotes the number of codebooks and W is the number of codewords in each codebook. w i = 1 rank ( i ) corresponds to the top-k list given from the teacher model. p q bk denotes the similarity score between the query q and the codeword c b k , i.e., p q bk = S ( q , c b k ) , and p d i bk denotes the similarity score between the candidate d i and the codeword c b k , i.e., p d i bk = S ( d i , c b k ) . The normalized value ˜ p q bk and ˜ p d i bk are calculated over the W codewords for each codebook through the softmax function.

This loss requires the enumeration of all the centroids, while the Weighted KL loss only includes parts of the centroids corresponding to the quantized function. It also aligns with the goal of nearest searching for the query with the learnable centroids as the bridge.

C.2 Experimental Performances

We compare the effectiveness of the three different distillation losses, i.e., Weighted KL loss, Distributed-based loss, and Lambdarank loss as reported in Table 8.

Findings. Overall, the Distributed-based loss leads to comparatively better performances than Weighted KL loss and Lambdarank loss. Compared with Weighted KL Loss and Lambdarank Loss, Distributed-based Loss gains the 2.03% and 3.54% improvements of Recall@10, 0.38%, and 0.70% of NDCG@10, respectively. The Lambdarank Loss concerns more about the relationships between the pair of items, while the other two care about the whole ranking order of the list. The weighted KL loss actually optimizes parts of the centroids, depending on which query vectors and candidate vectors are quantized, to match the ranking list, while all of the centroids are updated in the Distributed-based loss since the probabilities are calculated over all the centroids. Furthermore, the Distributed-based loss requires the similarity calculation between the original input vectors and the centroids, which eliminates the error caused by the compressed functions. The last observation is that Distributed-based Loss works better on inner product metric datasets, since it obtains the average improvements of 2.24% and 3.33% of Recall@10 on L2 distance and inner product, respectively, wherein the overall improvements for inner-product similarities.

D Search Efficiency and Retrieval Quality

The search efficiency and the comprise performance are important indicators of a compression search index. Thus, we report the effectivenesslatency performance for inference in Figure 4.

The time efficiency is measured by the latency of getting the Top-10 candidates for each query. The retrieval performance and the running time are governed by the number of posting lists we probe among the codebooks. The compressed ANNS methods (ScaNN and BLISS) and an uncompressed method (HNSW) are introduced for comparison. We can see that KDindex substantially outper-

Time-Recall trade-off.SIFT1M

Time-Recall trade-off. GIST1M

1.0

1.0 -

0.8

0.8

10

0.6

ScaNN

BLISS

HNSW

ngtqg

0.6

0.4

ScaNN

BLISS

0.4

qsgngt

vamana

0.2

HNSW

KDindex(AQ)

KDindex(AQ)

0.2

0.0

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

Time(s)

Time(s)

forms baselines in terms of both effectiveness and time efficiency. Specifically, we find that KDindex cost less time to arrive at a similar Recall than both ScaNN and BLISS. Although HNSW can achieve better results in the early stage, KDindex will surpass it after half a second.

E Limitations and Future Works

In this paper, we propose a novel knowledge distillation framework for high dimension index, which reduce storage obviously and can learn neighbor information from the teacher model. Especially, KDindex is independent to label (such as interaction information in the recommendation system or ground-truth neighbors in ANNS), which makes it flexible to be applied in more label-free scenarios. In the future, we will try more student models such as lattice quantization, whose codes already imply neighbors relationship. And we will take labels into account to improve retrieval performance progressively. We will further improve our work to benefit the broad community.

Peer reviews:
comments: SUMMARY:
The paper proposes a new method called KDindex to learn lightweight indexes by distilling knowledge from high-quality ANNS. The method outperforms existing learnable quantization-based indexes and state-of-the-art ANNS methods by learning to keep the ranking order, adding the reconstruction loss to minimize the compression error, and adopting a balanced partition strategy

STRENGTHS:
KDindex achieves good performance which outperforms existing learnable quantization-based indexes and some well-performed ANNS methods.

WEAKNESSES:
1.	This paper misses some related work, such as Poeem, JPQ, and MoPQ.
2.	The idea of this paper is similar with Distill-VQ. It improves the Distill-VQ by adding reconstruction loss into the learning objective and adjusts the posting list according to a balance strategy. I would like to see the comparison with Distill-VQ.
3.	For search efficiency and retrieval quality experiment, it does not compare with the state-of-the-art ANN algorithms in the ANN-Benchmarks, such as NGT-qg, qsgngt, and vamana.
4.	It seems balance strategy does not help too much.

QUESTIONS:
Please address the above weaknesses.

LIMITATIONS:
It does not include limitation analysis. Please add a Limitations section.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes to train an quantization-based approximate nearest neighbor search index using query and its approximate kNN items which are obtained using a teacher kNN search model. Authors propose loss functions to incorporate the ranking of kNN items as well as additional constraints to prevent trivial solutions and improve quality of approximation. The proposed framework yield significant improvement over baseline quantization methods as well as previous work on learning-based  quantization approaches.\n\nOverall score updated from `6: Weak Accept` to `7: Accept` after author response.\n\nSTRENGTHS:\n- The proposed method yields strong empirical results and experiment section includes ablation study, analysis of sensitivity to hyper-parameters, and comparison with state-of-the-art methods.\n- Each component of the proposed work is well-motivated and is accompanied with corresponding ablation results (although some of them have been moved to appendix but I would encourage authors to include some more details about baselines in the main paper).\n\nWEAKNESSES:\n- The presentation of the method and results can be improved further.\n    - For instance, a clear description of test-time inference process would be helpful.\n    - Also, it is not intuitively clear why using query-item interactions can results in better quantization for kNN search. The goal of the quantizer is to accurately express the original distance function. Does the proposed method improves kNN search performance at the expense of quantization performance or is the resulting quantization also a better approximation of the original distance function.\n- Missing information on Indexing time.\n    - While the proposed method does improve test-time performance, it would be interesting to see how much time does it take to index a given set of items\n- (Minor) Distillation is perhaps not the right term here.\n    - The model is trained using approximate nearest neighbors as per an expressive kNN search index.  But unlike “standard” distillation papers, the teacher model is not used to provide any soft training signal. The only training signal from query-item interactions is the ranking of items and this ranking is not induced by the teacher model but by the underlying similarity function.\n    - Why did the authors not used exact kNN items to train the index? I understand that obtaining exact kNN item might be significantly slower than retrieving approximate kNN items using teacher kNN index. But finding exact kNN items is a one-time cost and may turn out to be a small fraction of the overall running/training time as this exact kNN may not be the as expensive as the subsequent training of the kNN index.\n- (Minor) Proposed method requires a set of training queries to index the items. In typical settings for kNN search, the indexing algorithm has access to only the set of items. It would be interesting to use a subset of items as (pseudo-)queries in order to learn such an index.\n\nQUESTIONS:\n- What is the accuracy of the teacher model? Consider adding performance of teacher model in Table 1.\n- What is the intuition behind query-item interaction being helpful over just item-based quantizers?\n- Eq 6 is references in Algo 2 but I could not find it in the paper. Did authors mean eq 1?\n- Description of `w/o distillation` ablation says that “trains the encoder with knowledge distilled from the teacher model (HNSW in the experiment).” Can authors elaborate on this point in line 229? “*The improvement of KDindex is more significant when the distribution between query space and database space is different”.*\n    - What do authors mean in Line 231 by “The similarity could not  be obtained by original quantization methods”? Is this some fundamental limitation of the quantization-based methods?\n- Is it not clear how the proposed method avoid alternating between updating codewords and query/item-to-codeword assignments.\n    - In line 14 as Algo 2, as the codewords are updated, the index assignments in computed in Line 5 will become out-dated, right?\n    - Are the index assignments (assignments of documents to codewords) also updated somehow together with the codewords in line 14?\n\nLIMITATIONS:\nNo", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work addresses the problem of learning a lightweight index for high-dimensional similarity search problems. Lightweight index is desirable in many applications which can't afford a high precision heavyweight index due to higher storage cost or computational constraints. Unlike past work on lightweight index learning, this work assumes the availability of a heavyweight index at training time and explores the possibility of learning of the lightweight index with the guidance of the heavyweight index. It proposes a knowledge distillation framework for learning a lightweight index when no label information is available. The key idea is to use the top-K nearest neighbor results of every training query for guiding the learning of the lightweight index. This enables the use of ranking-oriented loss in the training of the  lightweight index. The work employs two tricks to avoid trivial and imbalanced solutions - i) a reconstruction loss that minimizes the distance between the query/candidate and its codeword ii) balancing of posting list.  The proposed framework is applied on four datasets and experimental results are discussed.\n\n\nFor every query, a cross-entropy like loss is computed with respect the top-K approximate nearest neighbors returned by the heavyweight teacher search model. This loss downweighs the similarity score between the codeword and the candidate with the reciprocal of the candidate's rank.\n\nSTRENGTHS:\nExperimental results show that knowledge distillation from a heavyweight index improves the retrieval performance of the lightweight indexes.\n\nAblation study shows that two of the three strategies employed make a significant difference to retrieval performance of the lightweight indexes.\n\nWEAKNESSES:\nAs the knowledge distillation process involves retrieving the top-K results for each query using the teacher search model and this is repeated in each iteration, the computational cost of distillation is high. \n\nDatasets used in the experiments are small in size (< 10M documents).\n\nNo discussion of and comparison with LSH and learning to hash techniques.\n\nQUESTIONS:\n\"Randomly initialize 133 the centroids (codewords) and assign indexes\" --> is it possible to take the help of teacher index to initialize in a more informed manner?\n\nIn Figure 2, what is the unit of time in x-axis?\n\nIn Figure 2, KDIndex stands for which specific student model? Is it PQ or OPQ or AQ? \n\nIn Figure 2, HNSW is time efficient compared to KDIndex for low to moderately high recall scenarios. However, in high recall scenario, HNSW is significantly worse than KDIndex. How can this happen as HNSW is the teacher and KDIndex is the student?\n\nIn Figure 2, best recall is achieved by ScaNN though this comes at increased search latency. However, KDIndex plateaus off quickly and increased search latency doesn't seem to help. \n\nThe recall numbers of KDIndex in Figure 2 for both SIFT1M and GIST1M datasets don't match with the recall numbers reported in Table 1. In fact Recall @10 in Table 1 for KDIndex is less than 35 for SIFT1M and less than 22 for GIST1M whereas Figure 2 reports Recall @10 as high as 0.8. What explains this discrepancy?\n\nTable 1 should report the results for the teacher search index HNSW to give an idea of the relative performance.\n\n In Section 4.4, it would be interesting to know the effect of B and W on storage and search latency.\n\nOn what criteria was B = 8  W = 256 was chosen as the optimal hyper-parameter setting? Was this the optimal setting for all the datasets?\n\nIn Table 2, it would be interesting to include K = 1 and 2. As K = 5 and K = 10 give very similar recall and MRR, it would be good to find out if smaller K also does similarly well. \n\nIn Section 4.5 and Table 2, what exactly is initialization warmed by quantization methods? Are you initializing the centroids with those for AQ/PQ/OPQ instead of random initialization?\n\nOf the three strategies employed by KDIndex, Balance seems to give least improvement in retrieval performance going b y Table 3 (for instance, PQ Recall @10 8.64 vs 8.62). However, Balance adds significant complexity to the training algorithm. It would be good to report and compare time taken for training KDIndex and w/o Balance in Table 3 to get a better understanding of the tradeoff between incremental improvement in retrieval performance and training complexity.\n\nIn Table 4, KDIndex refers to which of KDIndex(AQ), KDIndex(PQ), KDIndex(OPQ)? \n\nIn Table 4, why is compression for SIFT1M is much lower than GIST1M and other datasets (7 vs 63)?\n\nWhy are the similarity functions for MS MARCO different from that for SIFT1M and GIST1M?\n\nWhat is the additional time complexity of KDIndex relative to AQ, PQ and OPQ?\n\nWhy haven't 10M SIFT dataset and 5M SIFT dataset been used in the experiments as done by [38]?\n\nReferences [27] and [28] are one and the same!\n\nFigure 2: \"trage off\" should be \"trade-off\"\n\nLIMITATIONS:\nThe submission doesn't discuss the limitations of its work and the potential negative societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-09"}

comments: SUMMARY:
The paper proposes a method to compress indexes for ANNS by using knowledge distillation. They propose to use a  graph-based index teacher model and use the top-k nearest results obtained from the teacher indexes to act as the supervision signals to optimize the compressed function. The student model is optimized to have the same ranking orders as the teacher models. Different from previous work, they use a differentiable training process that updates the centroids and indexes simultaneously per mini-batch.

STRENGTHS:
Vector search is an important direction. This work tries to compress the embedding index, which is a critical task that can save huge storage and also increase the search performance. 

The results seem promising. The authors conduct experiments on 4 benchmarks and show that KDindex achieves a 40x index compression ratio, and 2x CPU speedup compared to the non-compressed method (HNSW).

WEAKNESSES:
There are several important questions that are not answered in this paper.

The paper doesn’t mention the training time for the KDindex. In addition to the storage, the indexing time for the index is also important. 

One other question is about how to index new documents. The student model training depends on a trained teacher model (HNSW). What if there is a new document added to the index corpus, will the distilled student generalize well for the new document? 

Also, it would be interesting to see larger benchmark, e.g. those have 1B entries, and see if the distillation would work.

QUESTIONS:
Why not reporting the HNSW performance in table 1?

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-14