SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Aniket Rege*† Aditya Kusupati*† Sharan Ranjit S† Alan Fan† Qingqing Cao†, Sham Kakade‡ Prateek Jain Ali Farhadi†
†University of Washington, ‡Google Research, ‡Harvard University
{kusupati,ali}@cs.washington.edu,prajain@google.com
Web-scale search systems learn an encoder to embed a given query which is then hooked into an approximate nearest neighbor search (ANNS) pipeline to retrieve similar data points. To accurately capture tail queries and data points, learned representations typically are rigid, high-dimensional vectors that are generally used as-is in the entire ANNS pipeline and can lead to computationally expensive retrieval. In this paper, we argue that instead of rigid representations, different stages of ANNS can leverage adaptive representations of varying capacities to achieve significantly better accuracy-compute trade-offs, i.e., stages of ANNS that can get away with more approximate computation should use a lower-capacity representation of the same data point. To this end, we introduce AdANNS [2], a novel ANNS design framework that explicitly leverages the flexibility of Matryoshka Representations [31]. We demonstrate state-of-the-art accuracy-compute trade-offs using novel AdANNS-based key ANNS building blocks like search data structures (AdANNS-IVF) and quantization (AdANNS-OPQ). For example on ImageNet retrieval, AdANNS-IVF is up to  $1.5\%$  more accurate than the rigid representations-based IVF [48] at the same compute budget; and matches accuracy while being up to  $90\times$  faster in wall-clock time. For Natural Questions, 32-byte AdANNS-OPQ matches the accuracy of the 64-byte OPQ baseline [13] constructed using rigid representations - same accuracy at half the cost! We further show that the gains from AdANNS translate to modern-day composite ANNS indices that combine search structures and quantization. Finally, we demonstrate that AdANNS can enable inference-time adaptivity for compute-aware search on ANNS indices built non-adaptively on matryoshka representations. Code is open-sourced at https://github.com/RAIVNLab/AdANNS.
Semantic search [24] on learned representations [40, 41, 50] is a major component in retrieval pipelines [4, 9]. In its simplest form, semantic search methods learn a neural network to embed queries as well as a large number  $(N)$  of data points in a  $d$ -dimensional vector space. For a given query, the nearest (in embedding space) point is retrieved using either an exact search or using approximate nearest neighbor search (ANNS) [21] which is now indispensable for real-time large-scale retrieval.
Existing semantic search methods learn fixed or rigid representations (RRs) which are used as is in all the stages of ANNS (data structures for data pruning and quantization for cheaper distance computation; see Section 2). That is, while ANNS indices allow a variety of parameters for searching the design space to optimize the accuracy-compute trade-off, the provided data dimensionality is typically assumed to be an immutable parameter. To make it concrete, let us consider inverted file index (IVF) [48], a popular web-scale ANNS technique [16]. IVF has two stages (Section 3) during inference: (a) cluster mapping: mapping the query to a cluster of data points [36], and (b) linear
(a) Image retrieval on ImageNet-1K.
(b) Passage retrieval on Natural Questions.
Figure 1: AdANNS helps design search data structures and quantization methods with better accuracy-compute trade-offs than the existing solutions. In particular, (a) AdANNS-IVF improves on standard IVF by up to  $1.5\%$  in accuracy while being  $90\times$  faster in deployment and (b) AdANNS-OPQ is as accurate as the baseline at half the cost! Rigid-IVF and Rigid-OPQ are standard techniques that are built on rigid representations (RRs) while AdANNS uses matryoshka representations (MRs) [31].
scan: distance computation w.r.t all points in the retrieved cluster to find the nearest neighbor (NN). Standard IVF utilizes the same high-dimensional RR for both phases, which can be sub-optimal.
Why the sub-optimality? Imagine one needs to partition a dataset into  $k$  clusters for IVF and the dimensionality of the data is  $d - \mathrm{IVF}$  uses full  $d$  representation to partition into  $k$  clusters. However, suppose we have an alternate approach that somehow projects the data in  $d/2$  dimensions and learns  $2k$  clusters. Note that the storage and computation to find the nearest cluster remains the same in both cases, i.e., when we have  $k$  clusters of  $d$  dimensions or  $2k$  clusters of  $d/2$  dimensions.  $2k$  clusters can provide significantly more refined partitioning, but the distances computed between queries and clusters could be significantly more inaccurate after projection to  $d/2$  dimensions.
So, if we can find a mechanism to obtain a  $d/2$ -dimensional representation of points that can accurately approximate the topology/distances of  $d$ -dimensional representation, then we can potentially build significantly better ANNS structure that utilizes different capacity representations for the cluster mapping and linear scan phases of IVF. But how do we find such adaptive representations? These desired adaptive representations should be cheap to obtain and still ensure distance preservation across dimensionality. Post-hoc dimensionality reduction techniques like SVD [14] and random projections [25] on high-dimensional RRs are potential candidates, but our experiments indicate that in practice they are highly inaccurate and do not preserve distances well enough (Figure 2).
Instead, we identify that the recently proposed Matryoshka Representations (MRs) [31] satisfy the specifications for adaptive representations. Matryoshka representations pack information in a hierarchical nested manner, i.e., the first  $m$ -dimensions of the  $d$ -dimensional MR form an accurate low-dimensional representation while being aware of the information in the higher dimensions. This allows us to deploy MRs in two major and novel ways as part of ANNS: (a) low-dimensional representations for accuracy-compute optimal clustering and quantization, and (b) high-dimensional representations for precise re-ranking when feasible.
To this effort, we introduce AdANNS, a novel design framework for semantic search that uses matryoshka representation-based adaptive representations across different stages of ANNS to ensure significantly better accuracy-compute trade-off than the state-of-the-art baselines.
Typical ANNS systems have two key components: (a) search data structure to store datapoints, (b) distance computation to map a given query to points in the data structure. Through AdANNS, we address both these components and significantly improve their performance. In particular, we first propose AdANNS-IVF (Section 4.1) which tackles the first component of ANNS systems. AdANNS-IVF uses standard full-precision computations but uses adaptive representations for different IVF stages. On ImageNet 1-NN image retrieval (Figure 1a), AdANNS-IVF is up to  $1.5\%$  more accurate for the compute budget and  $90\times$  cheaper in deployment for the same accuracy as IVF.
We then propose AdANNS-OPQ (Section 4.2) which addresses the second component by using AdANNS-based quantization (OPQ [13]) – here we use exhaustive search overall points. AdANNS-OPQ is as accurate as the baseline OPQ on RRs while being at least  $2 \times$  faster on Natural Questions [32] 1-NN passage retrieval (Figure 1b). Finally, we combine the two techniques to obtain AdANNS-IVFOPQ (Section 4.3) which is more accurate while being much cheaper – up to  $8 \times$  – than the traditional IVFOPQ [24] index. To demonstrate generality of our technique, we adapt AdANNS to DiskANN [22] which provides interesting accuracy-compute tradeoff; see Table 1.
While MR already has multi-granular representations, careful integration with ANNS building blocks is critical to obtain a practical method and is our main contribution. In fact, Kusupati et al. [31] proposed a simple adaptive retrieval setup that uses smaller-dimensional MR for shortlisting in retrieval followed by precise re-ranking with a higher-dimensional MR. Such techniques, unfortunately, cannot be scaled to industrial systems as they require forming a new index for every shortlisting provided by low-dimensional MR. Ensuring that the method aligns well with the modern-day ANNS pipelines is important as they already have mechanisms to handle real-world constraints like load-balancing [16] and random access from disk [22]. So, AdANNS is a step towards making the abstraction of adaptive search and retrieval feasible at the web-scale.
Through extensive experimentation, we also show that AdANNS generalizes across search data structures, distance approximations, modalities (text & image), and encoders (CNNs & Transformers) while still translating the theoretical gains to latency reductions in deployment. While we have mainly focused on IVF and OPQ-based ANNS in this work, AdANNS also blends well with other ANNS pipelines. We also show that AdANNS can enable compute-aware elastic search on prebuilt indices without making any modifications (Section 5.1); note that this is in contrast to AdANNS-IVF that builds the index explicitly utilizing "adaptivity" in representations. Finally, we provide an extensive analysis on the alignment of matryoshka representation for better semantic search (Section 5.2).
We make the following key contributions:
- We introduce AdANNS, a novel framework for semantic search that leverages matryoshka representations for designing ANNS systems with better accuracy-compute trade-offs.
- AdANNS powered search data structure (AdANNS-IVF) and quantization (AdANNS-OPQ) show a significant improvement in accuracy-compute tradeoff compared to existing solutions.
- AdANNS generalizes to modern-day composite ANNS indices and can also enable compute-aware elastic search during inference with no modifications.
Approximate nearest neighbour search (ANNS) is a paradigm to come as close as possible [7] to retrieving the "true" nearest neighbor (NN) without the exorbitant search costs associated with exhaustive search [21, 52]. The "approximate" nature comes from data pruning as well as the cheaper distance computation that enable real-time web-scale search. In its naive form, NN-search has a complexity of  $\mathcal{O}(dN)$ ;  $d$  is the data dimensionality used for distance computation and  $N$  is the size of the database. ANNS employs each of these approximations to reduce the linear dependence on the dimensionality (cheaper distance computation) and data points visited during search (data pruning).
Cheaper distance computation. From a bird's eye view, cheaper distance computation is always obtained through dimensionality reduction (quantization included). PCA and SVD [14, 26] can reduce dimensionality and preserve distances only to a limited extent without sacrificing accuracy. On the other hand, quantization-based techniques [6, 15] like (optimized) product quantization ((O)PQ) [13, 23] have proved extremely crucial for relatively accurate yet cheap distance computation and simultaneously reduce the memory overhead significantly. Another naive solution is to independently train the representation function with varying low-dimensional information bottlenecks [31] which is rarely used due to the costs of maintaining multiple models and databases.
Data pruning. Enabled by various data structures, data pruning reduces the number of data points visited as part of the search. This is often achieved through hashing [8, 46], trees [3, 12, 16, 48] and graphs [22, 38]. More recently there have been efforts towards end-to-end learning of the search data structures [17, 29, 30]. However, web-scale ANNS indices are often constructed on rigid  $d$ -dimensional real vectors using the aforementioned data structures that assist with the real-time search. For a more comprehensive review of ANNS structures please refer to [5, 34, 51].
Composite indices. ANNS pipelines often benefit from the complementary nature of various building blocks [24, 42]. In practice, often the data structures (coarse-quantizer) like IVF [48] and HNSW [37] are combined with cheaper distance alternatives like PQ [23] (fine-quantizer) for massive speed-ups in web-scale search. While the data structures are built on  $d$ -dimensional real vectors, past works consistently show that PQ can be safely used for distance computation during search time. As evident in modern web-scale ANNS systems like DiskANN [22], the data structures are built on  $d$ -dimensional real vectors but work with PQ vectors  $(32 - 64$ -byte) for fast distance computations.
ANNS benchmark datasets. Despite the Herculean advances in representation learning [19, 42], ANNS progress is often only benchmarked on fixed representation vectors provided for about a dozen million to billion scale datasets [1, 47] with limited access to the raw data. This resulted in the improvement of algorithmic design for rigid representations (RRs) that are often not specifically designed for search. All the existing ANNS methods work with the assumption of using the provided  $d$ -dimensional representation which might not be Pareto-optimal for the accuracy-compute trade-off in the first place. Note that the lack of raw-image and text-based benchmarks led us to using ImageNet-1K [45] (1.3M images, 50K queries) and Natural Questions [32] (21M passages, 3.6K queries) for experimentation. While not billion-scale, the results observed on ImageNet often translate to real-world progress [28], and Natural Questions is one of the largest question answering datasets benchmarked for dense passage retrieval [27], making our results generalizable and widely applicable.
In this paper, we investigate the utility of adaptive representations – embeddings of different dimensionalities having similar semantic information – in improving the design of ANNS algorithms. This helps in transitioning out of restricted construction and inference on rigid representations for ANNS. To this end, we extensively use Matryoshka Representations (MRs) [31] which have desired adaptive properties in-built. To the best of our knowledge, this is the first work that improves accuracy-compute trade-off in ANNS by leveraging adaptive representations on different phases of construction and inference for ANNS data structures.
3 Problem Setup, Notation, and Preliminaries
The problem setup of approximate nearest neighbor search (ANNS) [21] consists of a database of  $N$  data points,  $[x_1,x_2,\ldots ,x_N]$ , and a query,  $q$ , where the goal is to "approximately" retrieve the nearest data point to the query. Both the database and query are embedded to  $\mathbb{R}^d$  using a representation function  $\phi :\mathcal{X}\to \mathbb{R}^{d}$ , often a neural network that can be learned through various representation learning paradigms [2, 19, 20, 40, 42].
Matryoshka Representations (MRs). The  $d$ -dimensional representations from  $\phi$  can have a nested structure like Matryoshka Representations (MRs) [31] in-built -  $\phi^{\mathrm{MR}(d)}$ . Matryoshka Representation Learning (MRL) learns these nested representations with a simple strategy of optimizing the same training objective at varying dimensionalities. These granularities are ordered such that the lowest representation size forms a prefix for the higher-dimensional representations. So, high-dimensional MR inherently contains low-dimensional representations of varying granularities that can be accessed for free - first  $m$ -dimensions ( $m \in [d]$ ) ie.,  $\phi^{\mathrm{MR}(d)}[1:m]$  from the  $d$ -dimensional MR form an  $m$ -dimensional representation which is as accurate as its independently trained rigid representation (RR) counterpart -  $\phi^{\mathrm{RR}(m)}$ . Training an encoder with MRL does not involve any overhead or hyperparameter tuning and works seamlessly across modalities, training objectives, and architectures.
Inverted File Index (IVF). IVF [48] is an ANNS data structure used in web-scale search systems [16] owing to its simplicity, minimal compute overhead, and high accuracy. IVF construction involves clustering (coarse quantization through k-means) [36] on  $d$ -dimensional representation that results in an inverted file list [53] of all the data points in each cluster. During search,  $d$ -dimensional query representation is assigned to the most relevant cluster  $(C_i; i \in [k])$  by finding the closest centroid  $(\mu_i)$  using an appropriate distance metric  $(L_2$  or cosine). This is followed by an exhaustive linear search across all data points in the cluster which gives the closest NN (see Figure 5 in Appendix A for IVF overview). Lastly, IVF can scale to web-scale by utilizing a hierarchical IVF structure within each cluster [16]. Table 2 in Appendix A describes the retrieval formula for multiple variants of IVF.
Optimized Product Quantization (OPQ). Product Quantization (PQ) [23] works by splitting a  $d$ -dimensional real vector into  $m$  sub-vectors and quantizing each sub-vector with an independent  $2^b$
length codebook across the database. After PQ, each  $d$ -dimensional vector can be represented by a compact  $m \times b$  bit vector; we make each vector  $m$  bytes long by fixing  $b = 8$ . During search time, distance computation between the query vector and PQ database is extremely efficient with only  $m$  codebook lookups. The generality of PQ encompasses scalar/vector quantization [15, 36] as special cases. However, PQ can be further improved by rotating the  $d$ -dimensional space appropriately to maximize distance preservation after PQ. Optimized Product Quantization (OPQ) [13] achieves this by learning an orthonormal projection matrix  $R$  that rotates the  $d$ -dimensional space to be more amenable to PQ. OPQ shows consistent gains over PQ across a variety of ANNS tasks and has become the default choice in standard composite indices [22, 24].
Datasets. We evaluate the ANNS algorithms while changing the representations used for the search thus making it impossible to evaluate on the usual benchmarks [1]. Hence we experiment with two public datasets: (a) ImageNet-1K [45] dataset on the task of image retrieval – where the goal is to retrieve images from a database (1.3M image train set) belonging to the same class as the query image (50K image validation set) and (b) Natural Questions (NQ) [32] dataset on the task of question answering through dense passage retrieval – where the goal is to retrieve the relevant passage from a database (21M Wikipedia passages) for a query (3.6K questions).
Metrics Performance of ANNS is often measured using recall score [22],  $k$ -recall@  $N$  - recall of the exact NN across search complexities which denotes the recall of  $k$  "true" NN when  $N$  data points are retrieved. However, the presence of labels allows us to compute 1-NN (top-1) accuracy. Top-1 accuracy is a harder and more fine-grained metric that correlates well with typical retrieval metrics like recall and mean average precision (mAP@  $k$ ). Even though we report top-1 accuracy by default during experimentation, we discuss other metrics in Appendix C. Finally, we measure the compute overhead of ANNS using MFLOPS/query and also provide wall-clock times (see Appendix B.1).
Encoders. For ImageNet, we encode both the database and query set using a ResNet50  $(\phi_I)$  [19] trained on ImageNet-1K. For NQ, we encode both the passages in the database and the questions in the query set using a BERT-Base  $(\phi_N)$  [10] model fine-tuned on NQ for dense passage retrieval [27].
We use the trained ResNet50 models with varying representation sizes ( $d = [8, 16, \ldots, 2048]$ ; default being 2048) as suggested by Kusupati et al. [31] alongside the MRL-ResNet50 models trained with MRL for the same dimensionalities. The RR and MR models are trained to ensure the supervised one-vs-all classification accuracy across all data dimensionalities is nearly the same - 1-NN accuracy of 2048-d RR and MR models are  $71.19\%$  and  $70.97\%$  respectively on ImageNet-1K. Independently trained models,  $\phi_I^{\mathrm{RR}(d)}$ , output  $d = [8, 16, \ldots, 2048]$  dimensional RRs while a single MRL-ResNet50 model,  $\phi_I^{\mathrm{MR}(d)}$ , outputs a  $d = 2048$ -dimensional MR that contains all the 9 granularities.
We also train BERT-Base models in a similar vein as the aforementioned ResNet50 models. The key difference is that we take a pre-trained BERT-Base model and fine-tune on NQ as suggested by Karpukhin et al. [27] with varying (5) representation sizes (bottlenecks)  $(d = [48,96,\dots ,768]$ ; default being 768) to obtain  $\phi_N^{\mathrm{RR}(d)}$  that creates RRs for the NQ dataset. To get the MRL-BERT-Base model, we fine-tune a pre-trained BERT-Base encoder on the NQ train dataset using the MRL objective with the same granularities as RRs to obtain  $\phi_N^{\mathrm{MR}(d)}$  which contains all five granularities. Akin to ResNet50 models, the RR and MR BERT-Base models on NQ are built to have similar 1-NN accuracy for 768-d of  $52.2\%$  and  $51.5\%$  respectively. More implementation details can be found in Appendix B and additional experiment-specific information is provided at the appropriate places.
4 AdANNS - Adaptive ANNS
In this section, we present our proposed AdANNS framework that exploits the inherent flexibility of matryoshka representations to improve the accuracy-compute trade-off for semantic search components. Standard ANNS pipeline can be split into two key components: (a) search data structure that indexes and stores data points, (b) query-point computation method that outputs (approximate) distance between a given query and data point. For example, standard IVFOPQ [24] method uses an IVF structure to index points on full-precision vectors and then relies on OPQ for more efficient distance computation between the query and the data points during the linear scan.
Below, we show that AdANNS can be applied to both the above-mentioned ANNS components and provides significant gains on the computation-accuracy tradeoff curve. In particular, we present AdANNS-IVF which is AdANNS version of the standard IVF index structure [48], and the closely related ScaNN structure [16]. We also present AdANNS-OPQ which introduces representation adaptivity in the OPQ, an industry-default quantization. Then, in Section 4.3 we further demonstrate the combination of the two techniques to get AdANNS-IVFOPQ – an AdANNS version of IVFOPQ [24] – and AdANNS-DiskANN, a similar variant of DiskANN [22]. Overall, our experiments show that AdANNS-IVF is significantly more accuracy-compute optimal compared to the IVF indices built on RRs and AdANNS-OPQ is as accurate as the OPQ on RRs while being significantly cheaper.
4.1 AdANNS-IVF
Recall from Section 1 that IVF has a clustering and a linear scan phase, where both phase use same dimensional rigid representation. Now, AdANNS-IVF allows the clustering phase to use the first  $d_{c}$  dimensions of the given matryoshka representation (MR). Similarly, the linear scan within each cluster uses  $d_{s}$  dimensions, where again  $d_{s}$  represents top  $d_{s}$  coordinates from MR. Note that setting  $d_{c} = d_{s}$  results in non-adaptive regular IVF. Intuitively, we would set  $d_{c} \ll d_{s}$ , so that instead of clustering with a high-dimensional representation, we can approximate it accurately with a low-dimensional embedding of size  $d_{c}$  followed by a linear scan with a higher  $d_{s}$ -dimensional representation. Intuitively, this helps in the smooth search of design space for state-of-the-art accuracy-compute
Figure 2: 1-NN accuracy on ImageNet retrieval shows that AdANNS-IVF achieves near-optimal accuracy-compute trade-off compared across various rigid and adaptive baselines. Both adaptive variants of MR and RR significantly outperform their rigid counterparts (IVF-XX) while post-hoc compression on RR using SVD for adaptivity falls short.
trade-off. Furthermore, this can provide a precise operating point on accuracy-compute tradeoff curve which is critical in several practical settings.