Knowledge Distillation for High Dimensional Search Index
Authors: Zepu Lu
Jin Chen
Defu Lian
ZAIXI ZHANG
Yong Ge
Enhong Chen
Venue: NeurIPS 2023 poster (2023)
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
Zepu Lu Jin Chen', Defu Lian? Zhang??? Yong Ge", Enhong Chen 1.2
Sccooo of Computer Science and Technology, University of Science and Technology China State Key Laboratory of Cognitive Intelligence, Hefei, Anhui, China
University of Electronic Science and Technology of China *University of Arizona
zplu@ @plu@mail.ustcc.ammm mail. chenjin wtw.:sttt.dddd...ooeoeeoee liandefu wouttt.eeeee custc.ddu.cn, raix(@mail.sst..euuu... yongge aduu,na.edu, cheneh @ usstceed.cc
Abstract
Lightweight compressed indexes prevalent in Approximate Nearest Neighbor Search (ANNS) and Maximum Inner Product Search (MIPS) owing to their su- periority 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 high dimensions improve retrieval performance. In this paper, we propose a novel Knowledge Distillation 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 keep the same ranking order of the top-k relevant results yielded by the teacher model, which acts 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 re- sults demonstrate that KDindex outperforms existing learnable quantization- based indexes and is 40x lighter than the state-of-the-arr non-exhaustive methods while achieving comparable recall quality.
Introduction
Vector nearest neighbor search, which retrieves the most relevant vectors with the maximum similarity given the query vector, is fundamental task in information retrieval, such as image retrieval [13 (55. web search [5] 52| (53,, and item recommendation [32] 5 33
35 48). With significant number of candidate vectors and vector dimensions, the exact nearest neighbor search becomes intractable due the substantial computation costs and high query latency resulting in growing interest in Approximate Nearest Neighbors Search (ANNS).
Among the various indexes, the lightweight compressed indexes, especially those based quantiza- tion, are notable for their substantial advantages in terms of low storage costs and efficient parallel processing. As 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
"Corresponding Author
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

quantizers differ in whether there are explicit labels and whether they are related to queries. Tradi- tional clustering- methods for learning codebooks, e.g., PQ (24], OPQ [8] 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-basee 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-trutt -ouun labels to improve the retrieval performance, such the existing interactions between query and candidates LightRec [34], the ground-truth nearest neighbors in BLISS [2I]. 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 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 Knowledge Distillation for high dimensional search index framework (KDindex for short) to distill knowledge from 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 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 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 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 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 the supervision signals to guide the learning of centroids, which is label-free and takes the query information into consideration.
We utilize ranking-oriented oriented objectives as distillation loss with two learning constraints to avoid triv- ial solutions. Furthermore, we design the differentiable training process 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 -ort---rtttt-ttee compressed ANNS models.
2 Preliminaries
2.1 Problem Definition
Assume that there are M candidates {i,} such as documents, images, and items, in the retrieval system with no additional interaction information and the i-th candidate is represented by D- dimensional vector d, E D. D usually appears to huge numbers (larger than 100) nowadays retrieval systems. Given query vector g E DR.,, 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.

---

Input: (query)
Student Index
Step 1: Derive index assigoment query and documents via quantizer.
Query
Quantized
((d,,)
Codewords
Posting List
Balanced
Quantized
b. Teacher Index
Codeword
sinkhorn kkopp
((d,,)
Codeword
Codeword
Quantized Q(dx)
Distillation Loss
Top K Ranked documents list
top-K ranked candidates
Step 2: Update codewords and posting list by
DJ = (d,, 4,y. d..
Reconstrruttin loss and Distiliaionn loss.
Figure 1: Illustration of KDindex framework. Given the teacher index built on full-corpus, top-K ranked documents are returned for each query vector. Then, the query and corresponding documents compressed by lightweight indexes (such as PQ, OPQ, and AQ) and are ideally encoded with similar indexes, which are guided by the ranking-oriented distillation loss. Reconstruction loss and Sinkhorn-Know are used avoid trivial solutions.
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) widely used because of their strong performance in terms of both effectiveness and efficiency.
The quantization- -aaed indexes, a representative compressed ANNS algorithm, quantize each can- didate vector into the reconstruction vector from the B codebooks, each of which contains W vectors (codewords). We denote the b-th codebook as C and w-th codeword the b-th codebook as cy The Additive Quantization (AQ) encodes the vector as sum of B codewords, i.e., d Q(d) ΣB eb while the Product Quantization (PQ) and Optimized product quan- tization (OPQ) encode the vector as the concatenation of B codewords of each subspace, i.e., = (d) c + c² .... O c The codeword with the minimal distance from the original vector is chosen, i.e., cy (E(112....WW)
arg min ||æ di 11 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 also an effective method to compress high high-dimensional data into low-dimensional binarized codes, similar those clustering- based conventional quantization methods. Learning hash approaches have the property that objects that are close to each other have a higher probability of colliding than objects that far apart across various distance metries. 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 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, LightRee (34 exploits the interaction information between the certain query and document, and BLISS [2] takes the ground-truth neighbors over candidates to enhance the discrimination probability. These extra signals benefit better capturing the query- dependent 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 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- 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 high-precision teacher search model, which is well-leamed, and 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 {ilrank(ilg) < KJ. These candidates then behave weak signals to supervise the learning of codewords in the student models, where the student indexes are guided learn the same ranking orders and finally to return the similar indexes for the given query and candidate. Figure [ visualizes the whole training process of KDindex, with the three key components: Distillation Loss, two Learning Constraints, and 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 reconstruc- tion 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 7 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- loss as the distillation loss to guide the student indexes return the same nearest results. This objective function follows as:
B
W
(k. -- Σ
Ph log(Ps w))
&EDY b=1
where B denotes the number of codebooks and W is the number of codewords in each codebook. wi Fank(1) corresponds to the top-k list given from the teacher model. PhA denotes the similarity score between the query q and the codeword ct,, i.e., Ph S(q. ct,), ). and p denotes the similarity score between the candidate d, and the codeword ct. i.e., phr = ((d c²) The normalized value nd: ph. and Pbk 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, 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 imput 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:
Irecomstruccionn (x; C) D || Q(2)]1
where 1R 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 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 balanced clustering constraint to guide the candidate vectors be equally assigned to all quantization centroid embeddings. Here, we take one codebook as an example to formulate the learning objective with the constraint:
M W
min
- dill
W
M
subject to q(k|d,) {0, 1J,
(kidd,,) = 1, (k)dd.)
W
where qkk]d,) denotes the binary value indicating whether the candidate vector quantized the k-th centroid and the second condition ensures that only one centroid is selected. The last condition guides each centroid be allocated an equal number of candidates achieve balanced posting list.
Thus, the indexes for different candidates would be varied to avoid the trivial solution.
To solve Eq (3) relax the binary constraint qkk]d;) the continuous probability, i.e., (k:])d))) (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 algo- rithm 10 to solve it. Algorithm details the balance strategy to uniformly allocate the candidates.
Algorithm 1: Posting List Balance
Input: Document Vectors [d,) Codebook C = {ck W
Output:
Posting Lists (Px [P>
Initialize the posting lists with empty sets;
Generate the probability matrix P RRMx with the normalized similarity scores the ((d,,())
codewords
Σr exD S(dr,cd)
Get the transferred probability matrix ps E MxW by inputting P into Sinkhorn-Knopp: 4 for E (1,.., M) do
Determine the index according to the transferred probability, i.e., t arg max Pfw::
Add into the corresponding posting list P. P.U (i):
7 end
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 large deviation in the quantized vector and inaccuracy of the retrieval performance. To keep index assignment up-to-date during each training batch, we attempt to design a differentiable training manner for continuous computation. The Gumbel-Soffmax Trick 24 is common and popular solution, motivated by which we adopt similar sample operation to copy the gradient into the one-hot embeddings. Surprisingly, the probabilities of cach centroid are calculated in the
5

---

aforementioned balanced posting list and can be regarded as the weight of the centroids for updating. as shown Line If 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.
Algorithm 2: KDindex with Differentiable Training
Input: Training Queries Q = {qm Jm=1 Document Vectors Document Vectors {d. [d.)==1
Teacher Search Model T
Student Search Model including Codebooks and Index assignment Initialize the codebooks C:
Output:
2 while not reaching convergence conditions do
Get the quantized vector for the query ((q): 4 Sample mini-batch queries QB from Q: Assign indexes according to Algorithm
for q E 2B do
Retrieve the top-K nearest neighbors D& from the teacher search model T:
for E Di do
Get the one hot embedding of the index, i.c.,
c==One *;==OneHoo(t),& = arg max Pf2.. P.w):
Hot(t),
10
Copy the corresponding probability p; = [p11. Pi2: ..., piw];
11
Update the one hot encoding with gradient, i.e., 2, = (es - pi) detach( + pp
Get the quantized vector (d.)) é; C
13
end
14
Update the codebook C' depending on the distillation loss Eq and the reconstruction loss Eq (
15
end
16
end
3.6 Complexity Analysis
The training ways of KDindex and basic Qutization methods (AQ, PQ and OPQ) are different, therefore the complexity different 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 Table
Table 1: The time complexity 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 queries each batch), and the number of neighbors. N is the number of items. As for ScaNN, K, denotes the number of centroids in VQ vector quantization and Kp in PQ (
Product Quantization
Methods
Initialization
Training
Indexing O(NWD) B(NB(((D)B) + D2))
KDindex (PQ)
O(
OMMD + MBW + MBWK)
KDindex (OPQ)
O(MWD
O(MW + MBW + BBK)
KDindex (AQ)
O(MBWD)
O(MBWD D + MBW + MBWK)
O(N D)
HNSW
N/A
N/A
ONN N)
ScaNN
N/A
O(MWD MBW + MBWK)
(DN(BK, + K,D))
4 Experiment
4.1 Experiment Setup
Datasets and Metrics. Four large-scale retrieval benchmarks, including SIFTIM, GISTIM from
ANN datasets (2) .SS MARCO Doc and MS MARCO Passage from the TREC 2019 Deep Learning

---

Table 2: Comparison with Quantization- based methods. The improvement over the strongest baseline
statistically significant on a paired t-test (p < 0.05).
SIFTIM
GISTIM Recall 10
MS MARCO Doc
MS MARCO Passage Recall @ 10
Model
Recall @ 10
Recall 10
PQ [24]
31.
5.04=:.11
2.85+1032
1.80::0.08
OPQ
33.85:10.14
15.99=10.17
14.63-10.25
9.13:1..53
AQ 3
0..4
19.32=#019
16.73+0.55
10.18:0.69
DiffPQ
29.01
4.43=0022
4.17=0.26
2.53:0.06
DeepPQ 17
25.02=
4.39:10.18
6,47+10.3
8.43
GCD 25
33.58
15.76=0..13
15.59+0.27
10.28::046
RepCONC 53
33.59+0.32
15.32+00.66
16.07++0.31
10.37+0037
PQVAE 49
30.39+0.65
6.09++0.33
12.43++0.25
8.26++0.21
KDindex (PQ)
32.3
9.43
8.74
4.32
KDindex (OPQ)
34.7740.47
17.32=0.17
16.70+0.22
10.6610.69
KDindex (AQ)
37.30+0.11
21.33:1.23
18.9340.76
11.1940.26
NDCG@10 10
NDCG@10
MRR@ 10
MRR@ 10
PQ
73.21+0.17
19.74+0075
3.82++0.19
2.84
OPQ
75.76++0.09
49.10+0.74
34.75-+0.14
29.06
AQ
77.82+0.42
60.33+00.3
38.52++.008
33.52+
DiffPQ
0..53++0.16
17.39-00.64
10.43++0.399
3.69
DeepPQ
05.8
16.60=1.771
14.90110.18
28.3
GCD
75,4
48.82=0.58
38.8911.229
34.01
RepCONC
75.47
47.69-0.779
39.03-1.73
34.27
PQVAE
69.84
23.48-0.42
34.47
27.91
KDindex (PQ)
73.9
30.62
17.64
6.42
KDindex (OPQ)
76.32.40.63
52.36-10.49
39.75+0.25
34.62+0.47
KDindex (AQ)
80.07
63.17
41.69410.34
35.23
Track [9) used to validate the effectiveness of the proposed KDindex. SIFTIM and GISTIM both have exactly Million database points, and the dimension is 128 and 960, respectively. Document Retrieval consists of 3.2M documents, 0.36.M training queries, and 5K development queries. Passage Retrieval has a corpus of 8.8 M passages, 0.8 M training queries, and 0. M development queries. All the datasets offer the of training queries and testing queries. The detailed specifications of the datasets are shown in Table 6 The search model returns the top-K retrieval results given the test query and thus we use ranking-basee based metrics evaluate the performance, including Recall, NDCG and MRR with a cutoff of 10.
Implementation Details. SIFTIM and GISTIM datasets contain the candidate vectors with the dimension of 128 and 960 respectively. The vectors for MS MARCO Doc and MS MARCO Passage generated based huggingface transformers 477 and the dimension is previous works 14. The similarity function 5 for SIFTIM and GISTIM is the Euclidean distance and for MS MARCO Doc and Passage is the inner product. We exploit well-trained HNSW (37) model as the teacher model and K = 10 nearest neighbors are retrieved accordingly. We adopt PQ, OPQ, and AQ the student model for KDindex. For fair comparison, methods with optimized objectives proposed in KDindex instead of interaction information based loss function. Each vector quantized by B 8 codebooks, each of which contains W 256 codewords by default. The centroids are trained with learning rate of 0.01 and optimized by the Adam 28 optimizer. The batch size set to 64.
768 following
current supervised
4.2 Overall performance
We first compare the retrieval performance within the same retrieval time verify the effectiveness of KDindex. Experiments are run times to conduct the t-test and the performance of the retrieval is illustrated in Table from which we have the following findings:
DDindex respectively improves the retrieval performance of three native student models, demonstrat- ing the general effectiveness of KDindex. KDindex has relative 10.49% and 9.65% improvement in
7

---

Epoch Recoll SIFTIM PQ
Epoch Rocall SIFTIM
Epoch Recall SITTMM
32.5
37.0
Decal
04.23
24.00
35.5
82
102
40
60 Eacch
80
120
100
Eanch
Sanch
Spachs MS MARCO Doc PQ
Epoch Recall MS MARCO Dos OPQ
Epoch-Raceil MS MARCO Doc AQ
16.5
b5
17.5
17.0
150 375
130
100
Figure 2: Curves for recall during training warmed up by initialization.
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- -aaseddd-badddddee indexes with no ground 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
extremely difficult for the model to learn the codebooks as well as the index at the same time during the initialization phase in 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 differ- entiable training when the index assignment is approaching being balanced (maz) (P:l - min (Pyl < N W. j E W) where [Pa| 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 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 (Ai) 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 SIFTIM by KDindex(AQ)) L2 distance dataset and no more improvements can be obtained easily.
0.39% since the express ability of AQ strong on the
4.4 Memory Efficiency

---

As shown in Table 3 DDindex is more memory-effeient than the non-exhaustive ANNS method, HNSW, while performing com- parable effectiveness. The in- dex size of KDindex is less than 0.6GB, while HNSW re- quires storing all the original vectors and the graphs and con- sumes 0.5GB, 3.8GB, 9.2GB, and 25.3GB four datasets respectively for storing indexes. The average Compression Ratio between the memory of the non-exhaustive method (HNSW) and KDindex on four datasets ap- proximates 40x, demonstrating the superiority of the low memory cost of KDindex owing to the lightweight structure.
Table Memory (GB) of Indexes on different ANNS methods. Compression Ratio denotes the ratio between the memory of the non-exhaustive method (HNSW) and KDindex.
MS MARCO Doc
MS MARCO
Index(GBB)
SIFTIM
GISTIM
Passage
HNSW
0,51200 0.07036
3.84000
9. 9.19480 O. 1.19962
25.2967
KDindex
0.060644
0.53507
Compress Ratio
7x
63x
46x
46x
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 ex- periments on SIFTIM and compare the results of KDindex with different quantization methods.
PQ
DPO
AO
14.62
29
14.05
39.01
32.45
33.33
14.09
20.73
18.14
37
18.27
11.48
14.45
11.07
25.19
24
22
20.57
31.87
14.01
11.00
75-40
that, in terms of task performance, decreasing W has a far more traumatizing impact on KDin- dex. This is due to the nearest neigh- bor estimate losing accuracy when the number of codewords in each code- book drops.
Firstly,
30
23.64
40:51
40:73
22
35.72
16
10
16
12
16
10
34.24
29.00
34.79
21.30
34.23
10.48
34.30
40.33
42.92
20.00
50.90
43
23.39
34.73
42.23
43.21
22.87
37
4030
40.00
Secondly, increasing B or W would typically improve the task perfor- mance at the expense of lower effi- ciency and more storage cost, which means one can adjust B and W in the meantime to obtain trade-off be- tween effectiveness and efficiency.
20 81
73
60.79
45
22.55
$4.50
47.40
47.53
23.85
37
5093
51.10
Figure 3: Heat-maps of performance for various B and W values SIFTIM. The first line is the results of the origi- nal Quantization methods on Recall 10. The second line is KDindex based corresponding quantization methods. Darker is better.
Thirdly, notice that the small W and large B combination is preferable 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 SIFTIM.
Effects of K The relevant results come from the teacher model di- rectly influence the distribution be- tween the teacher and student and thus play vital role in distilla- tion loss. We vary the number of the nearest neighbors among
Table 4: The results under the different number of nearest neighbors in distillation loss on MS MARCO Doc.
K=5
K=10
K=20
K=50
K=100
Recall @ 10
18.32
18.93
17.56
17.18
16.43
MRR @ 10
41.32
41.69
40.33
39.44
39.03
{5, 10, 20, 50, 100} and report the results in Table
The model performs best when the number of neighbors is equal to 10, which consistent with the cutoff. When the number of neighbors increases. leading more data points are trained, the effeets

---

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) 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 the experiment). 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 Ginhhor-Knoop Knopp balance strategy. The reconstruction loss is so necessary to avoid trivial solution that could not be removed. (5) KDindex denotes methods that differentially train models with Reconstruction loss, Distillation loss, and Balance strategy.
Table 5: Ablation study on MSMARCO Doc dataset.
Model variants
Quantization
Initialization
Distillation loss
w/o Balance strategy
KDindex
@ 10
2.85
6.32
6.69
8.62
8.74
PQ
MRR@ 10
3.82
14.33
15.27
17.01
17.64
Recall@ 10
14.63
15.87
15.93
16.25
16.70
OPQ
MRR@ @ 10
34.75
38.96
39.01
39.22
39.75
Recall@ 10
16.73
17.39
17.60
18.26
18 93
AQ
MRR @ 10
38.52
36.52
37.02
40.39
41.69
We conduct experiments on the MS MARCO Doc dataset with the same settings mentioned above and report MRR 10 and Recall 10. Results are shown in Tab. We can see that all three strategies contribute to the effectiveness of KDindex. With the help of the distillation loss, Recall @ 10 and MRR @ achieve 14.34% and 10.01% % relative improvements respectively. This indicates that the knowledge distilled based 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 relative 2.61% 1% and 2.76% improvement 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 SinkhornKKoop 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 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-ae one.
Acknowledgments and Disclosure of Funding
The work was supported by grants from the National Natural Science Foundation of China (No. 60022222077 and 61976198).
10

---

References
[1] Martin Aumüller, Erik Bernhardsson, and Alexander Faithfull. Ann-benchmarks: A benchmark- ing 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.
[2] Martin Aumüller, Erik Bernhardsson, and Alexander Faithfull. ban-benchmarks: A bench- marking tool for approximate nearest neighbor algorithms. Information Systems, 87:101374, 2020.
[3] 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.
[4] Christopher Burges, Robert Ragno, and Quoe Le. Learning to rank with nonsmooth cost functions. Advances in neural information processing systems, 19, 2006.
[5] 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.
[6] Ting Chen, Lala Li, and Yizhou Sun. Differentiable product quantization for end-to-end embedding compression. In International Conference Machine Learning, pages 1617-1626. PMLR, 2020.
[7] Ting Chen, Martin Renqiang Min, and Yizhou Sun. Learning k-way d-dimensional discrete codes for compact embedding representations. In International Conference on Machine Learn- ing. pages 854-863. PMLR, 2018.
[8] 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
[9] Nick Craswell, Bhaskar Mitra, Emine Yilmaz, Daniel Campos, and Ellen M Voorhees. Overview of the trec 2019 deep learning track. arXiv preprint Stti:20033078220 2020.
[10] Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
[11] 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 Artificial Intelligence, volume 34, pages 51-58, 2020.
[12] Mayur Datar, Nicole Immorlica, Piotr Indyk, and Vahab S Mirrokni. Locality-sensitive sensitive hashing scheme based on p-stable distributions. Proceedings of the twentieth annual symposium on Computational geomelry, pages 253-262, 2004.
[13] 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
[14] Jacob Devlin, Ming- deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
Heig-WWiiee Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of
[15) 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.
[16] 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): 2023.
[17] Lianli Gao, Xiaosu Zhu, Jingkuan Song, Zhou Zhao, and Heng Tao Shen. Beyond product quan- tization: Deep progressive quantization for image retrieval. arXiv preprint arXiv: 1906.06698. 2019.
[18] 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.
[19] Ruiqi Guo, Sanjiv Kumar, Krzysztof Choromanski, and David Simeha. Quantization based fast inner product search. In Artificial intelligence and statistics, pages 482-490. PMLR, 2016.

---

[20] 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.
[21] Gaurav Gupta, Tharun Medini, Anshumali Shrivastava, and Alexander Smola. Bliss: A billion scale index using iterative re-partitioning. In Proceedings of the 28th ACM SICKDD Conference on Knowledge Discovery and Data Mining, pages 486-495, 2022.
[22] Geoffrey Hinton, Oriol Vinyals. Jeff Dean, et al. Distilling the knowledge in a neural network. arXiv preprint 1503.
2(7), 2015.
[23] Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterrization with gumbel-softmax. arXiv preprint 1611vv1611.01 144, 2016.
[24] 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.
[25] 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.cmm 220.000000,, 2022.
[26] Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion-scale similarity search with GPUs. IEEE Transactions on Big Data, 7(3):535-555 547, 2019.
[27] 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 614, 2020.
[28] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint urXiv:1412.6000 2014.
[29] 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.
[30] 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.
[31] Wen Li, Ying Zhang, Yifang Sun, Wei Wang, Mingjie Li, Wenjie Zhang, and Xuemin Lin. Approximate nearest neighbor search on high dimensional extteimeets,,ss,, analyses, and improvement. IEEE Transactions on Knowledge and Data Engineering. 4275):145--1488 2019.
[32] 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.
[33] Defu Lian, Qi Liu, and Enhong Chen. Personalized ranking with importance sampling. In Proceedings of The Web Conference 2020, pages 1093-1103, 2020.
[34] 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.
[35] Defu Lian, Yongji Wu, Yong Ge, Xing Xie, and Enhong Chen. Geography- awaree sequential location recommendation. In Proceedings of the 26th ACM SIGKDD international conference on knowledge discovery & data mining, pages 2009-2019, 2020.
[36] 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.
[37] Yu A Malkov and Dmitry A Yashunin. Efficient and robust approximate nearest neighbor search using hicrarchical navigable small world graphs. IEEE transactions on pattern analysis and machine intelligence, 42(4):824 836, 2018.
[38] Stanislav Morozov and Artem Babenko. Non-metric similarity graphs for maximum inner product search. Advances in Neural Information Processing Systems, 31, 2018.
[39] 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-2244 2014.
12

---

[40]
Behnam Neyshabur and Nathan Srebro. On symmetric and asymmetric Ishs for inner product search. In International Conference on Machine Learning. pages 1926-1934. PMLR, 2015.
[41] 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.
[42] Hanan Samet. Foundations of multidimensional and metric data structures. Morgan Kaufmann, 2006.
[43] Anshumali Shrivastava and Ping Li. Asymmetric Ish (alsh) for sublinear time maximum inner product search (mips). Advances in neural information processing systems, 27, 2014.
[44] Jiaxi Tang and Ke Wang. Ranking distillation: performance for recommender system. In Proceedings of the 24th ACM SIGKDD international conference on knowledge discovery & data mining, pages 2289-2298, 2018.
Learning compact ranking models with high
[45] Haoyu Wang, Defu Lian, and Yong Ge. Binarized collaborative filtering with distilling graph convolutional networks. arXiv preprint 190600//96 2019.
[46] 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.
[47] 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-thee art natural language processing. Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 38-45, Online, October 2020.
Association for Computational Linguistics.
[48] 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.
[49] Hanwei Wu and Markus Flierl. Learning product codebooks using vector-quantized autoen- coders for image retrieval. In 2019 IEEE Global Conference on Signal and Information Processing (GlobalSIP), pages 1-5. IEEE, 2019.
[50] 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.
[51] 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 Information Retrieval, pages 1513-1523, 2022.
[52] 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 Pro- ceedings of the 30th ACM International Conference on Information & Knowledge Management, pages 2487-2496, 2021.
[53] 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 Fifleenth ACM International Conference on Web Search and Data Mining, pages 1328-1336, 2022.
[54] Ting Zhang, Chao Du, and Jingdong Wang. Composite quantization for approximate nearest neighbor search. In International Conference Machine Learning, pages 838-846. PMLR, 2014.
[55] 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.
13

---

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, [3] and graph search [36 38 29 Vector compression
methods mainly aim to compress the index to accelerate retrieval. Popular algorithms include hashing and quantization
12
40
3 20 19]
Under the constraints of storage, compressed methods are widely investigated by researchers. Prod- uet quantization [24] / 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 B do not decompose space but directly learn multiple codebooks. There are also some algorithms that take query information into account. NEQ [TT] decomposes the quantization error into norm
error and direction error and improves existing VQ techniques for MIPS. ScaNN 22 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-baat baaedd conventional quantization methods. BLISS (1] 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 complex neural network was firstly trained and then transferred to 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, 4T have adopted knowledge
distillation to RS. RD 44 firstly proposes KD method that makes the student give high scores on
the top-ranked items of the teacher's 's recommendation list. Similarly, CD 33 makes the student imitate the teacher's '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 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 improve the performance of hashing- based [45] and quantization- [51) codes, The ranking information is distilled from graph-basce network enhance the performance of hashing codes. However, these works rely on the ground-trui labels (user-item interactions) to learn the ranking orders. This is different from our work, where label information 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 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.
14

---

Table 6: Statistics of the datasets.
Datasets
#Database
#Train
#Test
Dim
SIFTIM GISTIM
1,000,000
100,000
10,000
128
1,000,000
500,000
1,000
960
MS MARCO Doc
3,213,833
367,013
5,193
768
MS MARCO Passage
8,841,823
808,731
101,093
768
B More Details of Experimental Settings
B.1 Dataset Statisties
B.2
Baselines
The two groups of baseline ANNS models are compared to KDindex.
The first group Non-quantization-busell ANNS methods, which accelerate the search reducing the number of candidates. BLISS [2 adopts the learning-to-inder framework to learn the hashing-based compressed functions. ScaNN quantizes with anisotropic quantization loss functions which greatly penalizes the parallel component of a datapoint's 's residual relative to its orthogonal component. HNSW 37 builds hierarchical set of proximity graphs. Results can be found in Appendix D.
The second is Quantization-base ANNS methods, which compress the embeddings by hashing quantization functions. PQ (24] decomposes the vector representation space into the Cartesian product of subspaces. OPQ [8] jointly learns space decomposition and subspace quantization. AQ represents each vector as sum of several components each coming from separate codebook. The baselines are implemented based on the Faiss ANNS library [26]. The parameters B and W as KDiiddxx. DiffPQ 66, differentiable product quantization, generic and end-to-end learnable compression framework. DeepPQ (17), deep progressive quantization, end-to- end learns the quantization codes sequentially. PQ-VAE [49] 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 512- features. GCD learns rotation matrix via family of block Givens coordinate descent algorithms. RepCONC requires data points to uniformly clustered around the quantization centroids.
to be the
25
53
B.3
Implemental details
Table 7: Details of teacher model (HNSW).
Teacher (HNSW)
SIFTIM
GISTIM
MS MARCO Doc
MS MARCO Passage
32
32
32
32
efConstruction
40
40
100
100
efSearch
100
512
1024
1024
Search Time (s)
0.5862
1.3082
1,4805
4.7689
Building Time (s)
20.1s
2m25.4s
17m52s
98m17.2s
Recall @ 10
0.9865
0.9859
0.9292
0.9182
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 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 considered soft labels to supervise the learning
15

---

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 distill knowledge from the more accurate indexes to guide the student indexes to return the same nearest results.
The pair-wise ranking-basee based loss is widely used to learn the ranking list by leading the high-ranked candidate to have higher similarity scores. Lambdarank [ 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:
Lambdarank loss:
L(q, Di; = Σ
log (1 + exp(p; pÿ)) ANDCGG100]
i.JED&
where D'& denotes the top-k results retrieved from the teacher model and S(q. Q(d,)) is the similarity score between the query vector and the quantized vector of the candidate Q is the quantizer function related to the codebooks C. AN ANDCG& 10,, denotes the change with respect to NDCG@ 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-ranker items are more concerned. Finally, the loss function follows as:
Clg. Ck: = Σ
D log
P
(5)
where üt denotes the normalized ranked similarity score with the quantized vector and ji with the ground-truth vector. Specifically,
p: = w. Q(d.)) -= Sig,
Po and ji are the normalized values over the top-k retrieved candidates depending on the softmax function. w;; 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 learn better centroids. The introduced rank-oriented weight further guides the student index to return the same ranking list.
Distributed-base 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-basse loss function follows as:
B
W
Llg. =-
dd Por log(Psk w))
(6)
1EDY
b=
1=1
where B denotes the number of codebooks and W is the number of codewords in each codebook.
w; rank(i) corresponds to the top-k list given from the teacher model. p&a denotes the similarity score between the query q and the codeword et. i.e., Phs S(q, eg)), and po denotes the similarity score between the candidate d, and the codeword ct. i.e., di S(d, c The normalized value dd ph and Pbk 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 KI. loss, serributed-bassee loss, and Lambdarank loss reported in Table g
16

---

Findings. Overall, the Distributed-base loss leads to comparatively better performances than Weighted KL loss and Lambdarank loss. Compared with Weighted KL Loss and Lambdarank Loss, Distributed-basee 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-bass loss since the probabilities are calculated over all the centroids. Furthermore, the Distributed-baan basedd 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-bact based Loss works better inner product metric datasets, since it obtains the average improvements of 2.24% and 3.33% of Recall @ 10 on 1.2 distance and inner product, respectively, wherein the overall improvements for inner- inner-product similarities.
Table 8: The results of KDindex under different distillation loss functions.
SIFTIM
GISTIM
MS MARCO
MS MARCO Passage
Function
Recall 10
NDCG@ 10
Recall@ 10
NDCG@ 10
Recall@10
MRR@ 10
Rccall@
MRR@ 10
Lambdarank
36.32
20.69
62.06
11.10
34.86
Weighted KL Loss
36.68
79.33
21.02
62,75
18.24
40.93
11.06
34.63
Distributed hased Loss
37.30
80.01
21.33
63.17
18.93
41.69
11.19
35,23
D Search Efficiency and Retrieval Quality
The search efficiency and the com- prise performance are important in- dicators of compression search in- dex. Thus, report the effectiveness- latency performance for inference in Figure
Time Rocall trade off SFTIM
Recal bode GST.M
ScnV
LigS
9.65
The time efficiency is measured by the latency of getting the Top-10 can- didates for each query. The retrieval performance and the running time are governed by the number of post- ing lists probe among the code- books. The compressed ANNS meth- ods (ScaNN and BLISS) and an un- compressed method (HNSW) are in- troduced for comparison. We can see that KDindex substantially outper- 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.
(7e
Trabi
Figure 4: Comparison with more compressed ANNS models and uncompressed method (HNSW, ngtqg, qsgngt, vamana, and Milvus) in terms of effectiveness--atenny trade-off on two datasets. Up and left is better. The results of ScaNN, HNSW, ngtag, qsgngt, vamana, and Milvus come from ANN- Benchmarks
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 flexible to be applied in more label-free scenarios. In the future, we will try more student models such lattice quantization, whose codes already imply neighbors relationship. And we will take labels into account to improve retrieval performance progressively. We will further improve work to benefit the broad community.
17


Reviews:
[{"review_id": "zF92Ahi74y", "reviewer": "Reviewer_LzAv", "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.", "questions": "Why not reporting the HNSW performance in table 1?", "limitations": "N/A", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "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. \n\nThe 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.\n\nThe paper doesn’t mention the training time for the KDindex. In addition to the storage, the indexing time for the index is also important. \n\nOne 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? \n\nAlso, it would be interesting to see larger benchmark, e.g. those have 1B entries, and see if the distillation would work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "1.  [Training and Indexing time for the KDindex] \n\n    > **Comments**: 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.\n\n    Thank you for your suggestion, and we conduct an analysis on the complexity of training and indexing time.  \n    The whole training process includes the initialization with different conventional quantization-based methods, followed by the proposed learning-based approach. The various quantizers have an impact on the complexity analysis, and thus we discuss them separately.  \n    First of all, we describe the notations for clear understanding. 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 in the database. As for ScaNN, $K_v$ denotes the number of centroids in VQ ( vector quantization ) and $K_p$ in PQ (  Product Quantization ).  \n    In terms of the learning process, we analyze the time complexity of the forward training process during each batch, which encompasses quantization, balance strategy, and distillation loss calculation. The quantization complexity is intricately tied to the specific quantizers employed. The balance strategy exhibits a complexity of $O(MBW)$ and the distillation takes a complexity of $O(MBWK)$. We summarize the complexity with respect to each batch as Table 8 in the global response PDF.  \n\n    Considering the training process, the complexity of the balance strategy, i.e., $O(MBW)$, is considerably lower than that of quantization.  The number of neighbors often refers to small numbers, e.g., $K=10$, which are smaller than the dimension of embeddings ($D>100$). \n\n    [1] Sinkhorn distances: Lightspeed computation of optimal transport. NeurIPS 2013  \n2.  [how to index new documents] \n\n    > **Comments**: 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?\n\n    Given that our document corpus consists of over 1 million data points, the addition of a single new document is unlikely to significantly alter the overall data distribution. Unless the new document is notably distinct from the existing data, its influence on the distribution would be minor.  \n    As a result, the relationships between documents, both in the teacher index and the student model, tend to remain relatively stable even with the introduction of a new document. Thus, we can assign the new document with appropriate centroids according to the well-trained quantized index, i.e., $\\arg\\min_{i\\in1,2,\\dots, W}  \\| x - c_i^b \\| ^2$ where $i$ is the code (index) in the $b$-th codebook (subspace) for a new document $x$. The new document may share the same index with its similar document.  \n    More concretely, only a set of items are retrieved by the teacher index so that the learning of student index is accessible to the subset of items. The other items are assigned to certain index according to the well learned quantization-based index for later inference. The scale of the subset of items is closely related to the number of the neighbors $K$. We vary $K$in the experiment as shown in Table 2 to demonstrate the effectiveness of the well-learned KDindex with subset of items.\n\n3.  [Larger Dataset with 1B benchmark] \n\n    > **Comments**:  larger benchmark, e.g. those have 1B entries, and see if the distillation would work.\n\n    We add a 1B dataset Yandex Deep1B[1], which is a image descriptor dataset consisting of the projected and normalized outputs from the last fully-connected layer of the GoogLeNet model. The embeddings is pretrained on the Imagenet classification task. Due to the RAM resource constraints, we randomly sample 20M of the datasets.  The details are as follows:\n\n    | Datasets      | #Database  | #Train    | #Test  | Dim  |\n    | ------------- | ---------- | --------- | ------ | ---- |\n    | Yandex Deep1B | 20,000,000 | 7,000,000 | 10,000 | 96   |\n\n\n    \\#Train and \\#Test here represent the number of queries.\n\n    We compare the KDIndex with different quantizers in the following table, where KDindex outperforms all the original quantizers.\n\n    | Model     | PQ    | KDindex (PQ) | OPQ   | KDindex (OPQ) | AQ    | KDindex (AQ) |\n    | --------- | ----- | ------------ | ----- | ------------- | ----- | ------------ |\n    | Recall@10 | 9.61  | 12.47        | 16.84 | 18.77         | 17.81 | 18.57        |\n    | NDCG@10   | 33.85 | 37.71        | 45.34 | 49.32         | 53.63 | 55.15        |\n\n    [1] Efficient indexing of billion-scale datasets of deep descriptors. CVPR 2016 \n\n4. [Performance of HNSW]\n\n    > **Comments:** Why not reporting the HNSW performance in table 1?\n\n    For Table 1, we perform ADC over these quantization-based methods. Since HNSW is a graph-based method, the comparison with these quantization-based methods is not fair. We provide more details about training.\n\n    For training, We obtain the approximate top-K neighbors from teacher model and the performance is approaching Recall@10 0.99 by adding the search latency. The time and performance used in experimental are as follows: \n\n    | Datasets | SIFT1M | GIST1M | MS MARCO Doc | MS MARCO Passage |\n    | --- | --- | --- | --- | --- |\n    | Recall@10 | 0.9865 | 0.9859 | 0.9292 | 0.9182 |\n    | NDCG@10 | 0.9999 | 0.9999 | N/A | N/A |\n    | MRR@10 | N/A | N/A | 0.9493 | 0.9327 |\n    | Query time (s) | 0.5862 | 1.3082 | 1.4805 | 4.7689 |\n\n    For more comparison about the recall-time performance, please refer to Figure 6 in the global response PDF."}, {"review_id": "eU9HbO0X42", "reviewer": "Reviewer_uFQG", "summary": "This 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", "questions": "\"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\"", "limitations": "The submission doesn't discuss the limitations of its work and the potential negative societal impact of their work.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "Experimental 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", "weaknesses": "As 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "1.  [computational cost of distillation] \n\n> **Comments**: As 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\nThanks for your thoughtful approach to KDindex training. Initially, top-K results are fetched for all training queries after teacher index (HNSW) training, ensuring efficient retrieval. This step occurs once in real implementation, with results cached for future use, forming precomputed label generation. When reviewing the Algorithm, notably line 17, \"IO cost\" primarily matters. Cached results from file reduce computational load.\nAfter obtaining the top-K retrieved results, the time complexity of distillation loss calculation is $O(BWK)$ for each query. Compared with the quantization time, which takes $O(BWD)$, the distillation cost is comparable and acceptable. \n\n2.  [large datasets] Refer to reponse 3 for Reviewer LzAv.\n\n3.  [discussion and comparison with LSH and learning to hash techniques] \n\n> **Comments**: No discussion of and comparison with LSH and learning to hash techniques.\n\nThank you for your valuable suggestions. Actually, we have compared KDindex with the state-of-the-art hashing learning based method, i.e., BLISS [2], which preserves the ground-truth ranking orders to learn bucket partition for large-scale data. The experiments in section 4.3 demonstrate that our proposed KDindex outperforms BLISS in both retrieval performance and efficiency. \n\nFor more discussion about the hashing-based methods, they can be integrated into the discussion with quantization-based methods in introduction. LSH (Local sensitive hashing) is a data-independent un-supervised method, similar to those clustering-based conventional quantization methods. 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.  \n\nLearning to hash is also an effective method to compress high-dimensional data into low-dimensional binarized codes, where two types of loss functions are usually used. One is the reconstruction loss, which minimizes the distance between the original vector and the encoded vectors. The another one is the _ranking-based loss_, e.g., triplet loss and pair-wise loss, which encourages the model to learn a hash function that preserves the pairwise similarity relationships between positive points and negative points. However, only when the data set *contains the interaction information*, such as user clicks, the ranking loss would be utilized. In this paper, the training data has *no ground-truth labels (positive labels)*, so that the explicit ranking loss can not be adopted for optimization. The proposed KDindex incorporates extra supervision ranking-based signals through a teacher index to capture interactions between queries and items. \nThere are also numerous works that use knowledge distillation to improve the performance of hashing-based codes, such as [1], where 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. \n    \nFurthermore, our proposed framework can be adapted to a variety of compressed indexes. The learning to hash method would be instantiated as the student index in the subsequent works to show the generality of the model.\n\n[1] Binarized collaborative filtering with distilling graph convolutional networks. IJCAI 2019.  \n[2] BLISS: A Billion scale Index using Iterative Re-partitioning. SIGKDD 2022\n\n4. [The initialization of the codewords]\n\n> **Comments:** Is it possible to take the help of teacher index to initialize in a more informed manner?\n\nThank you for your valuable suggestion. A straightforward way is to exploit the neighbor information in the teacher index to aggregate the similar items with the same index. But for the centroids, the size and scale are remarkably different from that in HNSW. This would remain a further problem.\n\nKDindex is initialized with AQ/PQ/OPQ instead of random initialization. KDindex(AQ) represents that the quantization method is warmed up by AQ. The same goes for PQ and OPQ.  We also report that 'To accelerate the training, codebooks are warmed by original quantization methods such as PQ, OPQ, and AQ.' in Appendix D.\n\n5. [More descriptions abouth Figure 2]\n\nThe unit is second which corresponds with the results in [ann-benchmarks](https://ann-benchmarks.com/). KDindex here stands for AQ.\n\n> **Comments:** How can this happen as HNSW is the teacher and KDIndex is the student?\n\nHNSW's performance can be enhanced by extending search latency. In experiment, we take topk neighbor from HNSW whose recall@10 approaching 0.99 at the expense of search time. Thus, the teacher model provides more powerful signals for training student index.\n\n> **Comments:** 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\nWe increase the search latency for comparision, which can be referred in Figure 6 in the supplemented global response PDF. According to the figure, KDindex outperforms ScaNN and BLISS.\n\n6.  [More choices of K]\n\nWe conduct experiments on MS MARCO Doc datasets with $K=2$. The experiments are run for 5 times. The Recall@10 is $17.39\\pm0.25$ and MRR@10 is $40.10\\pm0.21$. Compared with more neighbors, e.g., $K=5$or $K=10$, extremely smaller number performs worse.\n"}, {"review_id": "UjKG8r1Dhw", "reviewer": "Reviewer_aAUh", "summary": "This 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.", "questions": "- 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?", "limitations": "No", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 2, "contribution": 3, "strengths": "- 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).", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "1.  [test-time inference process] \n\n     > **Comment**: a clear description of test-time inference process would be helpful.\n\n    For comparsion among quantization methods in Table 1,  Asymmetric distance computation (ADC) is used. \n    To compare with more ANNS methods in Table 2, an inverted file system is combined with the asymmetric distance computation (IVFADC). This allows rapid access to a small fraction of database indices and is shown to be successful for very large scale search. More detail can be found in [1]. \n    For KDindex, Table 1 is based on ADC and Figure 2 is based on IVFADC.  \n    [1] Product quantization for nearest neighbor search. TPAMI 2020.\n\n2. [Motivation of  query-item interactions]\n\n    > **Comment**: why using query-item interactions can results in better quantization for kNN search. \n\n    According to previous work[1], the introduction of query-item interactions would increase the probability of similar items assigning to the same centroids. Similarly, we extend the theoretical results to KDindex, which adopts the quantization methods rather than hashing based methods in [1]. \n    For the training query $q$ and its corresponding ground truth $p$ in the database, the expected probability of the centroids containing $p$ given the query $q$ increase by a positive margin after reassignment, i.e. $\\mathbb{E}[𝑓^′(e^′(𝒑)|𝒒)] ≥ \\mathbb{E}[𝑓 (e (𝒑)|𝒒)]$,where $f(\\cdot)$means the scoring function given by the model and $e(\\cdot)$means the quantization function that maps the points to centroids. The increment in this probability results in an increment in the quality of the retrieved candidates during inference.\n    The theorem implies that the centroids containing the relevant points derive a higher aggregated probability as it will contain other ground truths with higher probability. Consequently, query-item interaction is advantageous over item-based quantizers alone. The increment from interaction will be higher in the beginning as the relevant documents are quantized into the same centroids. As training goes on, the increment dies down to 0 and the model converges to the optimal quantization.\n\n    [1] BLISS: A Billion scale Index using Iterative Re-partitioning. SIGKDD 2022\n\n3. [Missing information on Indexing time.]\n\n    In the indexing stage, the codebooks are well learned and the indexing time of KDindex is same as corresponding basic quantization. We detail the time complexity in global response PDF.\n\n4.  [Distillation Details] \n\n    > **Comments:** Distillation is perhaps not the right term here\n\n    Conventional knowledge distillation fits in the classification tasks, where the more powerful teacher model imparts the soft labels to guide the learning of the student models. However, in the ranking tasks, such as recommendation and retrieval tasks, ranking distill methods are proposed, such as RD[1], RankDistill[2], to cater to the unique demands of ranking-oriented tasks. The ranking knowledge is distilled to the student model to learn a more accurate ranking performance. In this paper, KDindex adopts a more direct form of ranking signals from teacher index, where the approximate nearest neighbours are retrieved by the teacher index and the student models aim to preserve the ranking performance with the quantized embeddings. These retrieved topk items are different from different training queries and thus enhance the ability to capture similarities and conduct effective searches.  \n    [1] Ranking distillation: Learning compact ranking models with high performance for recommender system, KDD 2020  \n    [2] Rankdistil: Knowledge distillation for ranking, ICML 2021  \n\n\n5.  [Pseudo Query for training]\n\n    The large scale datasets with high dimensional vectors for ANN search usually include the training query vectors for learning and the test query vectors to fit in the dense retrieval scenarios. \n    For better illustrate the generalization ability, we conduct experiments with pseudo queries where we select 100K and 367K items as query vectors from the SIFT1M and MS MARCO Doc datasets. The results on SIFT1M of Pseudo query are approaching to the performance of True Query while on MS MARCO Doc is worse. The reason lies in the significantly different distribution between the query and database, as shown in Figure 5 of the supplemental PDF. Thus, distillation information from query vectors is helpful to learn codebooks.\n\n    |KDindex (AQ) | SIFT1M  |           |         | MS MARCO Doc |           |        |\n    | ------------ | ------- | --------- | ------- | ------------ | --------- | ------ |\n    |              | #query  | Recall@10 | NDCG@10 | #query       | Recall@10 | MRR@10 |\n    | True query   | 100,000 | 37.30     | 80.01   | 367,013      | 18.93     | 41.69  |\n    | Pseudo query | 100,000 | 37.15     | 79.43   | 367,013      | 16.64     | 38.25  |\n\n6. [Performance of Teacher Model] \n\n    For Table 1, we perform ADC over these quantization-based methods. Since HNSW is a graph-based method, the comparison with these quantization-based methods is not fair. We provide more details about training.\n\n    For training, We obtain the approximate top-K neighbors from teacher model and the performance is approaching ground-truth by adding the search latency. The time and performance used in experimental are as following: (More details in Appendix  B.3)\n\n\n    | Datasets       | SIFT1M | GIST1M | MS MARCO Doc | MS MARCO Passage |\n    | -------------- | ------ | ------ | ------------ | ---------------- |\n    | Recall@10      | 0.9865 | 0.9859 | 0.9292       | 0.9182           |\n    | NDCG@10        | 0.9999 | 0.9999 | N/A          | N/A              |\n    | MRR@10         | N/A    | N/A    | 0.9493       | 0.9327           |\n    | Query time (s) | 0.5862 | 1.3082 | 1.4805       | 4.7689           |\n\n\n    For more comparison about the recall-time performance, please refer to Figure 6 in the global response PDF.\n\n"}, {"review_id": "GYtQvsvC2i", "reviewer": "Reviewer_SPup", "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", "questions": "Please address the above weaknesses.", "limitations": "It does not include limitation analysis. Please add a Limitations section.", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "KDindex achieves good performance which outperforms existing learnable quantization-based indexes and some well-performed ANNS methods.", "weaknesses": "1.\tThis paper misses some related work, such as Poeem, JPQ, and MoPQ.\n2.\tThe 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.\n3.\tFor 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.\n4.\tIt seems balance strategy does not help too much.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "1.  [Other related works]\n\n    > **Comments:** This paper misses some related work, such as Poeem, JPQ, and MoPQ.\n\n    Thank you for your valuable advice. The mentioned related works, including Poeem, JPQ and MoPQ, focus more on the joint learning of both retrieval embedding models and the quantization models, where their aim is to integrate the separate encoding and compression process into an end-to-end training process. \n\n    Different from these works, our work mainly focuses on the better compressed methods given well-trained embedding vectors, where we \ndo not focus on the encoding models. We will elaborate on this difference in our work.\n\n2. [Discussion with Distill-VQ]\n\n    > **Comment:** 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.\n\n\n    Indeed, Distill-VQ incorporates the distillation method to jointly learn the query encoder and the index models. Distill-VQ can relax the requirements of ground-truth labels while KDindex is designed for unlabeled data. However, there are several differences between Distill-VQ and KDindex.\n\n    The first distinction lies in the dense embedding encoders. Distill-VQ necessitates the learning of query encoders for documents or  images to align with the candidate document embeddings. For KDindex, all the high dimensional embeddings are fixed so that it is a pure ANNs problem without the attention for encoding queries since we have no content information.\n\n    The second differentiation revolves around the nature of distillation signals. Distill-VQ takes the similarity information of fixed document embeddings as signals by disclosing the scores between the original fixed embeddings and the reconstructed embeddings. However, KDindex aims to preserve the ranking information from a teacher index model, where the teacher index model directly provides the retrieved top-k results. Thus, the sampled candidate document set plays an important role in Distill-VQ, which just represents local similarity information in the case of a huge document size.\n\n    For better compare the performance of different distillation methods, we conduct experiments for better illustration. For fair comparison, all the embeddings are fixed in Distill-VQ and the index structure follows AQ.  We adopt the in-batch sampling method with a batch size of 64 in Distill-VQ and calculate the distillation loss with  KL-divergence loss function. The following table shows the results:\n\n    | Models       | SIFT1M    |         | MS MARCO Doc |        |\n    | ------------ | --------- | ------- | ------------ | ------ |\n    |              | Recall@10 | NDCG@10 | Recall@10    | MRR@10 |\n    | Distill-VQ   | 35.79     | 78.43   | 17.34        | 39.93  |\n    | KDindex (AQ) | 37.30     | 80.01   | 18.93        | 41.69  |\n\n    As shown in the above table, KDindex still outperforms Distil-VQ, demonstrating the superiority of the distillation methods in KDindex.\n\n3. [Graph-based ANN Benchmarks]\n\n    > **Comments:** 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.\n\n\n    The results of NGT-qg, qsgngt, Milvus and vamana are added in Fig 2, which can be referred in Figure 6 in the supplement PDF. These alogrithms are graph based and full percision vectors are stored on node of graph. These graph-based methods show better performances especially with more latency. All these graph-based method can act as teacher model to train our KDindex.\n    We point that our KDindex is designed to enhance the performance of compressed index with the help of better-performed teacher index rather than defeat those graph-based methods.  As for compressed methods, such as ScaNN and BLISS, KDindex can perform better. In the futrue work, we will try more efficient method as teacher index to further improve compressed index. \n\n\n\n4. [Balance Strategy]\n\n    > **Comments:**  It seems balance strategy does not help too much. \n\n     Compared with other important parts, the relative 2.61% and 2.76% improvements in Recall@10 and MRR@10 of posting balance strategy are smaller but it still accounts for significant retrieval improvements. "}]

Author rebuttal:
Dear reviewers,

Thank you for all the time and effort you spent in writing reviews for KDindex. We polish the description and conduct more experiments folloing reviews. The figures (datasets distribution and time-recall curves) and table (time complexity) are attached in PDF.

Authors.

Meta-review: The paper proposes a method to compress indexes for ANNS by using knowledge distillation. All the reviewers seem excited about the proposal and there was a very healthy discussions where most concerns from the reviewers were addressed.  Overall, all the reviewers agree that the proposed contribution is worthy of publication.
Decision: Accept (poster)