SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zepu Lu $^{1,2}$ , Jin Chen $^{3}$ , Defu Lian $^{1,2*}$ , Zaixi Zhang $^{1,2}$ , Yong Ge $^{4}$ , Enhong Chen $^{1,2}$
$^{1}$ School of Computer Science and Technology, University of Science and Technology of China
$^{2}$ State Key Laboratory of Cognitive Intelligence, Hefei, Anhui, China
<sup>3</sup>University of Electronic Science and Technology of China
<sup>4</sup>University of Arizona
zplu@mail.ustc.edu.cn, chenjin@std.uestc.edu.cn, liandefu@ustc.edu.cn, zaixi@mail.ustc.edu.cn,
yongge@arizona.edu, cheneh@ustc.edu.cn
Lightweight compressed indexes are prevalent in Approximate Nearest Neighbor Search (ANNS) and Maximum Inner Product Search (MIPS) owing to their superiority of retrieval efficiency in large-scale datasets. However, results given by compressed indexes are less accurate due to the curse of dimension and limitation of optimization objectives (e.g., lacking interactions between queries and documents). Thus, we are encouraged to design a new learning algorithm for the compressed search index in high dimensions to improve retrieval performance. In this paper, we propose a novel 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 to keep the same ranking order of the top-k relevant results yielded by the teacher model, which acts as the additional supervision signals between queries and documents to learn the similarities between documents. Furthermore, to avoid the trivial solutions that all candidates are partitioned to the same post list, the reconstruction loss that minimizes the compressed error, and the posting list balance strategy that equally allocates the candidates, are integrated into the learning objective. Experiment results demonstrate that KDindex outperforms existing learnable quantization-based indexes and is  $40 \times$  lighter than the state-of-the-art non-exhaustive methods while achieving comparable recall quality.
Vector nearest neighbor search, which retrieves the most relevant vectors with the maximum similarity given the query vector, is a fundamental task in information retrieval, such as image retrieval [13, 55], web search [51, 52, 53], and item recommendation [32, 5, 33, 35, 48]. With a significant number of candidate vectors and vector dimensions, the exact nearest neighbor search becomes intractable due to the substantial computation costs and high query latency [42], resulting in a growing interest in Approximate Nearest Neighbors Search (ANNS).
Among the various indexes, the lightweight compressed indexes, especially those based on quantization, are notable for their substantial advantages in terms of low storage costs and efficient parallel processing. As a result, these indexes draw the attention from both academics and industries and are present in well-known open repositories such as FAISS [26]. Existing learning methods for
quantizers differ in whether there are explicit labels and whether they are related to queries. Traditional clustering-based methods for learning codebooks, e.g., PQ [24], OPQ [18], and AQ [3], are categorized as un-supervised algorithms only connected with document embeddings without the query information. Another type of technique [20] involves training the codebooks by optimizing the reconstruction-based loss, which relies on correlations with queries and documents but lacks explicit natural supervision signals. The last category of algorithms employs the query-dependent ground-truth labels to improve the retrieval performance, such as the existing interactions between query and candidates in LightRec [34], the ground-truth nearest neighbors in BLISS [21]. These additionally introduced supervision signals help capture the relationships between items and queries, thus contributing to better quantizers. However, the interaction label is only available in a small set of datasets and the expense of obtaining the ground-truth labels is particularly high, making the scenarios without any label information a more common case.
In this paper, we creatively propose a Knowledge Distillation for high dimensional search index framework (KDindex for short) to distill knowledge from a more advanced teacher index model into lightweight indexes, under the circumstances without available label information. The advanced indexes, such as graph-based indexes, can yield more accurate results owing to their powerful expressiveness of the vector space and hence becomes able to teach the less-accurate lightweight indexes. Specifically, the top-k nearest results obtained from the teacher indexes act as the supervision signals to optimize the compressed functions. Therefore, index distillation is formulated as the top-k learning for the student index model from the teacher index. Specific ranking-oriented losses are exploited to learn the knowledge from the teacher model, which guides the results from the student model to have the same ranking orders as the teacher models. To avoid the trivial solutions, where the candidate vectors collapse to the identical centroid, we have two learning constraints for the learning objective. The reconstruction loss minimizes the distance between the original input and the compressed input, which pushes the centroids by the different query and candidate vectors. The second strategy is the posting list balance, which encourages an equal distribution of documents among each centroid. Thus, the candidates would be equally assigned to different centroids to avoid trivial solutions. Lastly, unlike recent iterative approaches that compulsively modify indexes, our approach uses a differentiable training process that updates the centroids and indexes simultaneously per mini-batch. This aims to minimize the error resulting from asynchronous updates. In this way, the student model would yield better performance thanks to the powerful expressiveness of the neighbor relationships of the teacher model, while keeping the low inference time and storage cost.
The contribution of this work can be summarized as follows:
- To the best of our knowledge, this is the first attempt to extract knowledge from the high-precision search index with advanced structures into lightweight indexes in order to improve the retrieval performance of the lightweight index in high dimensions.
- We propose the index distillation paradigm to learn the top-k nearest neighbors retrieved from the teacher models, where the top-k retrieved results act as the supervision signals to guide the learning of centroids, which is label-free and takes the query information into consideration.
- We utilize ranking-oriented objectives as distillation loss with two learning constraints to avoid trivial solutions. Furthermore, we design the differentiable training process to avoid the asynchronous update of centroids and index assignments.
- Experimental results on four benchmark datasets demonstrate that KDindex achieves a 40x index compression ratio, and 2x CPU speedup while maintaining comparable retrieval performance as the state-of-the-art compressed ANNS models.
Assume that there are  $M$  candidates  $\{i_j\}_{j=1}^M$ , such as documents, images, and items, in the retrieval system with no additional interaction information and the  $i$ -th candidate is represented by a  $D$ -dimensional vector  $\pmb{d}_i \in \mathbb{R}^D$ .  $D$  usually appears to be huge numbers (larger than 100) in nowadays retrieval systems. Given a query vector  $\pmb{q} \in \mathbb{R}^D$ , the system attempts to return the top- $k$  relevant candidates depending on the similarity scores  $S(\cdot, \cdot)$  between the query and candidate vectors, such as the inner product, L2 distance, and cosine distance.
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 are 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-Knopp are used to 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) are widely used because of their strong performance in terms of both effectiveness and efficiency.
The quantization-based indexes, a representative compressed ANNS algorithm, quantize each candidate vector into the reconstruction vector from the  $B$  codebooks, each of which contains  $W$  vectors (codewords). We denote the  $b$ -th codebook as  $C^b$  and the  $w$ -th codeword in the  $b$ -th codebook as  $c_w^b$ . The Additive Quantization (AQ) encodes the vector as a sum of  $B$  codewords, i.e.,  $\tilde{d} = Q(d) = \sum_{b=1}^{B} c_{w_b}^b$ , while the Product Quantization (PQ) and Optimized product quantization (OPQ) encode the vector as the concatenation of  $B$  codewords of each subspace, i.e.,  $\tilde{d} = Q(d) = c_{w_1}^1 \oplus c_{w_2}^2 \oplus \ldots \oplus c_{w_B}^B$ . The codeword with the minimal distance from the original vector is chosen, i.e.,  $c_w^b = \arg \min_{i \in \{1, 2, \ldots, W\}} \| x - c_{w_i}^b \|^2$ .  $Q(\cdot)$  denotes the quantizer function. The codewords are usually learned by the unsupervised clustering algorithms or the reconstruction loss, and the posting lists which record the belonging candidates are then constructed. During the inference stage, the similarity scores between the query vector and the codewords are calculated and the approximate nearest candidates are oriented. Consequently, given the well-trained indexes, the overall inference time complexity is only correlated with the number of codewords and codebooks.
Learning to hash is also an effective method to compress high-dimensional data into low-dimensional binarized codes, similar to those clustering-based conventional quantization methods. Learning to hash approaches have the property that objects that are close to each other have a higher probability of colliding than objects that are far apart across various distance metrics. The drawbacks of these approaches are the requirement for a large number of hash tables in order to achieve good search quality and these methods are unmindful of the distribution of vectors, often leading to lop-sided partitions and long query times. Thus, we focus more on the quantization-based compressed methods.
3 Index Distillation
Despite the efficiency advantages of lightweight quantization-based indexes, they are still limited by the representation space of codewords. Existing works introduce additional ground-truth labels to supervise the learning. For example, LightRec [34] exploits the interaction information between the certain query and document, and BLISS [21] takes the ground-truth neighbors over candidates to enhance the discrimination probability. These extra signals benefit better capturing the 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 to learn the centroids. These works provide less relevant candidates, thus resulting in less accurate inference results. Intuitively, if we use more progressive index structures, such as graph-based methods, to encode the similarity information over the whole vector space, more accurate results would be obtained. This motivates the progressive indexes as the teacher to distill more informative high-dimensional knowledge to the lightweight quantization-based indexes, with the goal of improving the accuracy of lightweight indexes under high dimensions.
3.1 Overview of Index Distillation
The whole framework, i.e., Knowledge Distillation for high dimensional search index (KDindex), consists of a high-precision teacher search model, which is well-learned, and a lightweight student search model to be learned, where the student model attempts to learn knowledge from the top-k relevant candidates retrieved from the teacher model. Specifically, given the well-trained candidate vectors, when a query vector  $\mathbf{q}$  requests, the teacher search model returns the set of the top-k approximate nearest neighbors  $\mathcal{D}_K^T = \{i| \text{rank}(i|\mathbf{q}) \leq K\}$ . These candidates then behave as weak signals to supervise the learning of codewords in the student models, where the student indexes are guided to learn the same ranking orders and finally to return the similar indexes for the given query and candidate. Figure 1 visualizes the whole training process of KDindex, with the three key components: Distillation Loss, two Learning Constraints, and a Differentiable Training paradigm which are detailed in the following sections. The whole training process is illustrated in Algorithm 2.
3.2 Initialization
The initialization of the student search model involves the following steps: (1) Initialize the centroids (codewords) and assign indexes (2) Update the centroids based on the distillation loss and reconstruction loss given the training queries (3) Adjust the posting lists depending on the latest centroids. A straightforward process is to update the posting lists per  $T$  epochs after centroids updating is finished by enumerating all training queries, similar to the paradigm of BLISS [21].
3.3 Distillation Loss
Knowledge distillation was first proposed for classification tasks [22], where the probabilities of each class attained from the large-scale teacher network are considered as soft labels to supervise the learning of the small-size student network. The cross-entropy loss is commonly used as the distillation loss to minimize the difference between the teacher and student networks. Here, the teacher search model provides the top-k relevant candidates rather than the continuous value of probabilities. We exploit the following Distributed-based loss as the distillation loss to guide the student indexes to return the same nearest results. This objective function follows as:
$$
\mathcal {L} (\boldsymbol {q}, \mathcal {D} _ {K} ^ {T}; \boldsymbol {C}) = - \sum_ {i \in \mathcal {D} _ {K} ^ {T}} \sum_ {b = 1} ^ {B} \sum_ {k = 1} ^ {W} \tilde {p} _ {b k} ^ {q} \log \left(\tilde {p} _ {b k} ^ {d _ {i}} \cdot w _ {i}\right) \tag {1}
$$
where  $B$  denotes the number of codebooks and  $W$  is the number of codewords in each codebook.  $w_{i} = \frac{1}{\text{rank}(i)}$  corresponds to the top-k list given from the teacher model.  $p_{bk}^{q}$  denotes the similarity score between the query  $q$  and the codeword  $c_{k}^{b}$ , i.e.,  $p_{bk}^{q} = S(q, c_{k}^{b})$ , and  $p_{bk}^{d_{i}}$  denotes the similarity score between the candidate  $d_{i}$  and the codeword  $c_{k}^{b}$ , i.e.,  $p_{bk}^{d_{i}} = S(d_{i}, c_{k}^{b})$ . The normalized value  $\tilde{p}_{bk}^{q}$  and  $\tilde{p}_{bk}^{d_{i}}$  are calculated over the  $W$  codewords for each codebook through the softmax function.
This loss attempts to minimize the distance between the queries and top-k neighbors by calculating the similarity scores with all the centroids. Thus, we could obtain more information from centroids and focus on the top-K nearest neighbors.
3.4 Learning Constraints
Although the distillation loss supervises the learning of quantizers, the student model would collapse with a trivial solution where all candidates are assigned to the identical centroid. To avoid trivial solutions, we propose two strategies and we will introduce them one by one.
Reconstruction Loss. A straightforward solution is to design the reconstruction loss to draw the distance between the input vector and the reconstructed vector encoded through the quantized function. Specifically, MSE loss is adopted here to specialize the reconstruction loss for the both query and candidate vectors:
$$
\mathcal {L} _ {\text {r e c o n s t r u c t i o n}} (\boldsymbol {x}; \boldsymbol {C}) = \frac {1}{D} \| \boldsymbol {x} - Q (\boldsymbol {x}) \| ^ {2} \tag {2}
$$
where  $\pmb{x} \in \mathbb{R}^D$  represents the original input vector such as the query vector  $\pmb{q}$  and the candidate vector  $\pmb{d}$ .  $Q$  denotes the quantized function to reconstruct the input vector and the corresponding codewords would be updated through the loss.
Balanced Posting List. The posting list refers to the candidate list according to the certain centroids. The trivial solutions would result in an imbalanced posting list, where some posting lists would be substantially longer than others and thus resulting linear complexity in the worst case. Thus, we impose a balanced clustering constraint to guide the candidate vectors to be equally assigned to all quantization centroid embeddings. Here, we take one codebook as an example to formulate the learning objective with the constraint:
$$
\min  _ {\boldsymbol {C}} \sum_ {i = 1} ^ {M} \sum_ {k = 1} ^ {W} q (k | \boldsymbol {d} _ {i}) \| \boldsymbol {c} _ {w} - \boldsymbol {d} _ {i} \| ^ {2}
$$
$$
\text {s u b j e c t} q (k | \boldsymbol {d} _ {i}) \in \{0, 1 \}, \sum_ {k = 1} ^ {W} q (k | \boldsymbol {d} _ {i}) = 1, \sum_ {i = 1} ^ {M} q (k | \boldsymbol {d} _ {i}) = \frac {M}{W} \tag {3}
$$
where  $q(k|d_i)$  denotes the binary value indicating whether the candidate vector is quantized to the  $k$ -th centroid and the second condition ensures that only one centroid is selected. The last condition guides each centroid to be allocated an equal number of candidates to achieve a balanced posting list. Thus, the indexes for different candidates would be varied to avoid the trivial solution.
To solve Eq (3), we relax the binary constraint  $q(k|d_i)$  to the continuous probability, i.e.,  $q(k|d_i) \in (0,1)$ , to get the approximated solution. The problem then can be reduced to the optimal transport problem by taking the distance as the cost of quantization and we use the Sinkhorn-Knopp algorithm [10] to solve it. Algorithm 1 details the balance strategy to uniformly allocate the candidates.
Algorithm 1: Posting List Balance
Input: Document Vectors  $\{d_i\}_{i = 1}^M$  , Codebook  $C = \{c_k\}_{k = 1}^W$
Output: Posting Lists  $\{\mathcal{P}_k\}_{k = 1}^W$
1 Initialize the posting lists with empty sets ;
2 Generate the probability matrix  $\pmb{P} \in \mathbb{R}^{M \times W}$  with the normalized similarity scores over the codewords  $p_{ik} = \frac{\exp S(\pmb{d}_i, \pmb{c}_k)}{\sum_{k'=1}^W \exp S(\pmb{d}_i, \pmb{c}_{k'})}$ ;
3 Get the transferred probability matrix  $\pmb{P}^{s} \in \mathbb{R}^{M \times W}$  by inputting  $\pmb{P}$  into Sinkhorn-Knopp;
4 for  $i\in \{1,\dots,M\}$  do
5 Determine the index according to the transferred probability, i.e.,  $t = \arg \max (p_{i1}^{s}, p_{i2}^{s}, \dots, p_{iW}^{s})$ ;
6 Add  $i$  into the corresponding posting list  $\mathcal{P}_t = \mathcal{P}_t\cup \{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 a large deviation in the quantized vector and inaccuracy of the retrieval performance. To keep the index assignment up-to-date during each training batch, we attempt to design a differentiable training manner for continuous computation. The Gumbel-Softmax Trick [23] is a common and popular solution, motivated by which we adopt a similar sample operation to copy the gradient into the one-hot embeddings. Surprisingly, the probabilities of each centroid are calculated in the
aforementioned balanced posting list and can be regarded as the weight of the centroids for updating, as shown in Line 9-11 of Algorithm 2.
The whole training process is illustrated in Algorithm 2, which can be computed in parallel for each mini-batch. Accordingly, the learning of centroids and the index re-assignment are updated within the same mini-batch, rather than the asynchronous epochs.
Algorithm 2: KDindex with Differentiable Training
Input: Training Queries  $\mathcal{Q} = \{\pmb {q}_m\}_{m = 1}^N$  Document Vectors Document Vectors  $\{\pmb {d}_i\}_{i = 1}^M$  Teacher Search Model  $T$    
Output: Student Search Model including Codebooks and Index assignment   
1 Initialize the codebooks  $C$  .   
2 while not reaching convergence conditions do   
3 Get the quantized vector for the query  $Q(q)$  .   
4 Sample a mini-batch queries  $\mathcal{Q}_B$  from  $\mathcal{Q}$  .   
5 Assign indexes according to Algorithm 1;   
6 for  $q\in \mathcal{Q}_B$  do   
7 Retrieve the top-  $K$  nearest neighbors  $\mathcal{D}_K^T$  from the teacher search model  $T$  .   
8 for  $i\in \mathcal{D}_K^T$  do   
9 Get the one hot embedding of the index, i.e.,  $\pmb {e}_i = \mathrm{One\_Hot}(t),t = \arg \max (p_{i1}^s,p_{i2}^s,\dots,p_{iW}^s)$  . Copy the corresponding probability  $\pmb {p_i} = [p_{i1},p_{i2},\dots,p_{iW}]$  . Update the one hot encoding with gradient, i.e.,  $\tilde{e}_i = (\pmb {e}_i - \pmb {p}_i).detach() + \pmb {p}_i$  . Get the quantized vector  $Q(\pmb {d}_i)\approx \tilde{e}_i\cdot \pmb{C}$  .   
13 end   
14 Update the codebook  $C$  depending on the distillation loss Eq (6) and the reconstruction loss Eq (2);   
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 is different in the training phase. The indexing and inference way of KDindex and Qutization methods are consistent and there is no additional time complexity in the index and inference stages. More details are described in Table 1.
Table 1: The time complexity of KDindex, HNSW, and ScaNN. Denoted by  $D$  item embedding dimension,  $B$  the number of subspaces,  $W$  the number of centroids in each codebook,  $M$  the batch size (the number of queries in each batch), and  $K$  the number of neighbors.  $N$  is the number of items. As for ScaNN,  $K_v$  denotes the number of centroids in VQ (vector quantization) and  $K_p$  in PQ (Product Quantization).
Methods	Initialization	Training	Indexing
KDindex (PQ)	O(MWD)	O(MWD + MBW + MBWK)	O(NWD)
KDindex (OPQ)	O(MWD2)	O(MWD2 + MBW + MBWK)	O(NBW((D/B) + D2))
KDindex (AQ)	O(MBWD)	O(MBWD + MBW + MBWK)	O(NBWD)
HNSW	N/A	N/A	O(ND log N)
ScaNN	N/A	O(MWD + MBW + MBWK)	O(N(BKp(D/B) + KvD))
Datasets and Metrics. Four large-scale retrieval benchmarks, including SIFT1M, GIST1M from ANN datasets [2], MS MARCO Doc and MS MARCO Passage from the TREC 2019 Deep Learning