L EARNING S PACE P ARTITIONS FOR N EAREST N EIGHBOR S EARCH
Microsoft Piotr Indyk MIT
Ilya Razenshteyn Microsoft Research Tal Wagner MIT
A BSTRACT
Space partitions of R d underlie a vast and important class of fast nearest neighbor search (NNS) algorithms.
Inspired by recent theoretical work on NNS for general metric spaces (Andoni et al., 2018b;c), we develop a new framework for building space partitions reducing the problem to balanced graph partitioning followed by supervised classiﬁcation.
We instantiate this general approach with the
KaHIP graph partitioner (Sanders & Schulz, 2013) and neural networks, respectively, to obtain a new partitioning procedure called Neural Locality-Sensitive Hashing (Neural LSH).
On several standard benchmarks for NNS (Aum¨uller et al., 2017), our experiments show that the partitions obtained by
Neural LSH consistently outperform partitions found by quantization-based and tree-based methods as well as classic, data-oblivious LSH.
Introduction
The Nearest Neighbor Search (NNS) problem is deﬁned as follows.
Given an n -point dataset P in a d -dimensional
Euclidean space R d , we would like to preprocess P to answer k -nearest neighbor queries quickly.
That is, given a query point q ∈ R d , we want to ﬁnd the k data points from P that are closest to q .
NNS is a cornerstone of the modern data analysis and, at the same time, a fundamental geometric data structure problem that led to many exciting theoretical developments over the past decades.
See, e.g., Wang et al.
; Andoni et al.
(2018a)
for an overview.
The main two approaches to constructing efﬁcient NNS data structures are indexing and sketching .
The goal of indexing is to construct a data structure that, given a query point, produces a small subset of P (called candidate set ) that includes the desired neighbors.
Such a data structure can be stored on a single machine, or (if the data set is very large) distributed among multiple machines.
In contrast, the goal of sketching is to compute compressed representations of points to enable computing approximate distances quickly (e.g., compact binary hash codes with the Hamming distance used as an estimator, see the surveys Wang et al. Indexing and sketching can be (and often are) combined to maximize the overall performance (Wu et al., 2017; Johnson et al., 2017).
Both indexing and sketching have been the topic of a vast amount of theoretical and empirical literature.
In this work, we consider the indexing problem.
In particular, we focus on indexing based on space partitions .
The overarching idea is to build a partition of the ambient space R d and split the dataset P accordingly.
Given a query point q , we identify the bin containing q and form the resulting list of candidates from the data points residing in the same bin (or, to boost the accuracy, nearby bins as well).
Some of the popular space partitioning methods include locality-sensitive hashing (LSH) (Lv et al., 2007;
Andoni et al., 2015; Dasgupta et al., 2017); quantization-based approaches, where partitions are obtained via k -means clustering of the dataset (J´egou et al., 2011; Babenko & Lempitsky, 2012); and tree-based methods such as random-projection trees or PCA trees (Sproull, 1991; Bawa et al., 2005; Dasgupta & Sinha, 2013;
Keivani & Sinha, 2018).
Compared to other indexing methods, space partitions have multiple beneﬁts.
First, they are naturally applicable in distributed settings, as different bins can be stored on different machines (Bahmani et al., 2012;
Ni et al., 2017; Li et al., 2017; Bhaskara & Wijewardena, 2018).
Moverover, the computational efﬁciency of search can be further improved by using any nearest neighbor search algorithm locally on each machine.
Second, partition-based indexing is particularly suitable for GPUs due to the simple and predictable memory access pattern (Johnson et al., 2017).
Finally, partitions can be combined with cryptographic techniques to yield efﬁcient secure similarity search algorithms (Chen et al., 2019).
Thus, in this paper we focus on designing space partitions that optimize the trade-off between their key metrics: the number of reported candidates, the fraction of the true nearest neighbors among the candidates, the number of bins, and the computational efﬁciency of the point location.
Recently, there has been a large body of work that studies how modern machine learning techniques (such as neural networks) can help tackle various classic algorithmic problems (a partial list includes Mousavi et al.
; Baldassarre et al.
; Bora et al.
; Dai et al.
Metzler et al.
; Kraska et al.
; Balcan et al.
Lykouris & Vassilvitskii (2018); Mitzenmacher (2018); Purohit et al.
Similar methods—under the name “learn to hash”—have been used to improve the sketching approach to NNS (Wang et al., 2016).
However, when it comes to indexing , while some unsupervised techniques such as PCA or k -means have been successfully applied, the full power of modern tools like neural networks has not yet been harnessed.
This state of affairs naturally leads to the following general question: Can we employ modern (supervised) machine learning techniques to ﬁnd good space partitions for nearest neighbor search?
1. Our contribution
In this paper we address the aforementioned challenge and present a new framework for ﬁnding high-quality space partitions of R d .
Our approach consists of three major steps:
Build the k -NN graph G of the dataset by connecting each data point to k nearest neighbors; 2.
Find a balanced partition P of the graph G into m parts of nearly-equal size
such that
the number of edges between different parts is as small as possible; 3.
Obtain a partition of R d by training a classiﬁer on the data points with labels being the parts of the partition P found in the second step.
See Figure 1 for illustration.
The new algorithm directly optimizes the performance of the partition-based nearest neighbor data structure.
Indeed, if a query is chosen as a uniformly random data point, then the average k -NN accuracy is exactly equal to the fraction of edges of the k -NN graph G whose endpoints are separated by the partition P .
This generalizes to out-of-sample queries provided that the query and dataset distributions are close, and the test accuracy of the trained classiﬁer is high.
At the same time, our approach is directly related to and inspired by recent theoretical work (Andoni et al., 2018b;c) on
NNS for general metric spaces.
In particular, using the framework of (Andoni et al., 2018b;c), we prove that, under mild conditions on the dataset P , the k -NN graph of P can be partitioned with a hyperplane into two parts of comparable size such that only few edges get split by the hyperplane.
This gives a partial theoretical justiﬁcation of our method.
The new framework is very ﬂexible and uses partitioning and learning in a black-box way.
This allows us to plug various models (linear models, neural networks, etc.) and explore the trade-off between the quality and the algorithmic efﬁciency of the resulting partitions.
We emphasize the importance of balanced partitions for the indexing problem, where all bins contain roughly the same number of data points.
This property is crucial in the distributed setting, since we naturally would like to assign a similar number of points to each machine.
Furthermore, balanced partitions allow tighter control of the number of candidates simply by varying the number of retrieved parts.
Note that a priori, it is unclear how to partition R d so as to induce balanced bins of a given dataset.
Here the combinatorial portion of our approach is particularly useful, as balanced graph partitioning is a well-studied problem, and our supervised extension to R d naturally preserves the balance by virtue of attaining high training accuracy.
We speculate that the new method might be potentially useful for solving the NNS problem for non-Euclidean metrics, such as the edit distance (Zhang & Zhang, 2017) or optimal transport distance (Kusner et al., 2015).
Indeed, for any metric space, one can compute the k -NN graph and then partition it.
The only step that needs to be adjusted to the speciﬁc metric at hand is the learning step.
Let us ﬁnally put forward the challenge of scaling our method up to billion-sized or even larger datasets.
For such scale, one needs to build an approximate k -NN graph as well as using graph partitioning algorithms that are faster than KaHIP.
We leave this exciting direction to future work.
For the current experiments (datasets of size 10 6 points), preprocessing takes several hours.
Another important challenge is to obtain NNS algorithms based on the above partitioning with provable guarantees in terms of approximation and running time.
However, we expect it to be difﬁcult, in particular, since all the current state-of-the-art NNS algorithms lack such guarantees (e.g., k -means-based (J´egou et al., 2011) or graph methods (Malkov & Yashunin, 2018), see also (Aum ¨ uller et al., 2017) for a recent SOTA survey).
Evaluation We instantiate our framework with the KaHIP algorithm (Sanders & Schulz, 2013) for the graph partitioning step, and either linear models or small-size neural networks for the learning step.
We evaluate it on several standard benchmarks for NNS (Aum¨uller et al., 2017) and conclude that in terms of quality of the resulting partitions, it consistently outperforms quantization-based and tree-based partitioning procedures, while maintaining comparable
[Figure]
algorithmic efﬁciency.
In the high accuracy regime, our framework yields partitions that lead to processing up to 2 .
3 × fewer candidates than the strongest baseline.
As a baseline method we use k -means clustering (J ´ egou et al., 2011).
It produces a partition of the dataset into k bins, in a way that naturally extends to all of R d , by assigning a query point q to its nearest centroid.
(More generally, for multi-probe querying, we can rank the bins by the distance of their centroids to q ).
This simple scheme yields very high-quality results for indexing.
Besides k -means, we evaluate LSH (Andoni et al., 2015), ITQ (Gong et al., 2013),
PCA tree (Sproull, 1991), RP tree (Dasgupta & Sinha, 2013), and Neural Catalyzer (Sablayrolles et al., 2019).
1. Related work
On the empirical side, currently the fastest indexing techniques for the NNS problem are graph-based (Malkov &
Yashunin, 2018).
The high-level idea is to construct a graph on the dataset (it can be the k -NN graph, but other constructions are also possible), and then for each query perform a walk, which eventually converges to the nearest neighbor.
Although very fast, graph-based approaches have suboptimal “locality of reference”, which makes them less suitable for several modern architectures.
For instance, this is the case when the algorithm is run on a GPU (Johnson et al., 2017), or when the data is stored in external memory (Sun et al., 2014) or in a distributed manner (Bahmani et al., 2012; Ni et al., 2017).
Moreover, graph-based indexing requires many rounds of adaptive access to the dataset, whereas partition-based indexing accesses the dataset in one shot.
This is crucial, for example, for nearest neighbor search over encrypted data (Chen et al., 2019).
These beneﬁts justify further study of partition-based methods.
Machine learning techniques are particularly useful for the sketching approach, leading to a vast body of research under the label “learning to hash”
(Wang et al., 2014; 2016).
In particular, several recent works employed neural networks to obtain high-quality sketches (Liong et al., 2015; Sablayrolles et al., 2019).
The fundamental difference from our work is that sketching is designed to speed up linear scans over the dataset, by reducing the cost of distance evaluation, while indexing is designed for sublinear time searches, by reducing the number of distance evaluations.
We note that while sketches are not designed for indexing, they can be used for that purpose, since a b -bit hashing scheme induces a partition of R d into 2 b parts.
Nonetheless, our experiments show that partitions induced by these methods (such as Iterative Quantization (Gong et al., 2013)) are not well-suited for indexing, and underperform compared to quantization-based indexing, as well as to our methods.
We highlight in particular the recent work of Sablayrolles et al.
(2019), which uses neural networks to learn a mapping f : R d → R d (cid:48) that improves the geometry of the dataset and the queries to facilitate subsequent sketching.
It is natural to ask whether the same family of maps can be applied to enhance the quality of partitions for indexing.
However, as our experiments show, in the high accuracy regime the maps learned using the algorithm of Sablayrolles et al.
(2019) consistently degrade the quality of partitions.
Finally, we mention that here is some prior work on learning space partitions: Cayton & Dasgupta (2007); Ram & Gray (2013);
Li et al.
However, all these algorithms learn hyperplane partitions into two parts (then applying them recursively).
Our method, on the other hand, is much more ﬂexible, since neural networks allow us to learn a much richer class of partitions.
Our method
Given a dataset P ⊆ R d of n points, and a number of bins m > 0 , our goal is to ﬁnd a partition R of R d into m bins with the following properties:
Balanced:
The number of data points in each bin is not much larger than n/m . 2.
Locality sensitive:
For a typical query point q ∈ R d , most of its nearest neighbors belong to the same bin of R .
We assume that queries and data points come from similar distributions.
Simple: The partition should admit a compact description and, moreover, the point location process should be computationally efﬁcient.
For example, we might look for a space partition induced by hyperplanes.
Formally, we want the partition R that minimizes the loss E q (cid:104)(cid:80)
[Equation]
η )( n/m ) , where q is sampled from the query distribution, N k ( q ) ⊂ P is the set of its k nearest neighbors in P , η >is a balance parameter, and R ( p ) denotes the part of R that contains p .
First, suppose that the query is chosen as a uniformly random data point , q ∼ P .
Let G be the k -NN graph of P , whose vertices are the data points, and each vertex is connected to its k nearest neighbors.
Then the above problem boils down to partitioning vertices of the graph G into m bins such that each bin contains roughly n/m vertices, and the number of edges crossing between different bins is as small as possible (see Figure 1(b)).
This balanced graph partitioning problem is extremely well-studied, and there are available combinatorial partitioning solvers that produce very high-quality solutions.
In our implementation, we use the open-source solver KaHIP (Sanders & Schulz, 2013), which is based on a sophisticated local search.
More generally, we need to handle out-of-sample queries, i.e., which are not contained in P .
Let (cid:101) R denote the partition of G (equivalently, of the dataset P ) found by the graph partitioner.
To convert (cid:101) R into a solution to our problem, we need to extend it to a partition R of the whole space
R d that would work well for query points.
In order to accomplish this, we train a model that, given a query point q ∈ R d , predicts which of the m bins of (cid:101) R the point q belongs to (see
Figure 1(c)).
We use the dataset P as a training set, and the partition (cid:101) R as the labels – i.e., each data point is labeled with the ID of the bin of (cid:101) R containing it.
The method is summarized in Algorithm 1.
The geometric intuition for this learning step is that – even though the partition (cid:101) R is obtained by combinatorial means, and in principle might consist of ill-behaved subsets of R d – in most practical scenarios, we actually expect it to be close to being induced by a simple partition of the ambient space.
For example, if the dataset is fairly well-distributed on the unit sphere, and the number of bins is m = 2 , a balanced cut of G should be close to a hyperplane.
The choice of model to train depends on the desired properties of the output partition R .
For instance, if we are interested in a hyperplane partition, we can train a linear model using SVM or regression.
In this paper, we instantiate the learning step with both linear models and small-sized neural networks.
Here, there is natural tension between the size of the model we train and the accuracy of the resulting classiﬁer, and hence the quality of the partition we produce.
A larger model yields better NNS accuracy, at the expense of computational efﬁciency.
We discuss this in Section 3.
Multi-probe querying Given a query point q , the trained model can be used to assign it to a bin of a partition R , and search for nearest neighbors within the data points in that part.
In order to achieve high search accuracy, we actually train the model to predict several bins for a given query point, which are likely to contain nearest neighbors.
For neural networks, this can be done naturally by taking several largest outputs of the last layer.
By searching through more bins (in the order of preference predicted by the model) we can achieve better accuracy, allowing for a trade-off between computational resources and accuracy.
Hierarchical partitions When the required number of bins m is large, in order to improve the efﬁciency of the resulting partition, it pays off to produce it in a hierarchical manner.
Namely, we ﬁrst ﬁnd a partition of R d into m 1 bins, then recursively partition each of the bins into m 2 bins, and so on, repeating the partitioning for L levels.
The total number of bins in the overall partition is m = m 1 · m 2 · . . .
m L .
See Figure 2 for illustration.
The advantage of such a hierarchical partition is that it is much simpler to navigate than a one-shot partition with m bins.
Neural LSH with soft labels In the primary instantiation of our framework, we set the supervised learning component to a a neural network with a small number of layers and constrained hidden dimensions (the exact parameters are speciﬁed in the next section).
In order to support effective multi-probe querying, we need to infer not just the bin that contains the query point, but rather a distribution over bins that are likely to contain this point and its neighbors.
A
T -probe candidate list is then formed from all data points in the T most likely bins.
In order to accomplish this, we use soft labels for data points generated as follows.
For S ≥ 1 and a data point p , the soft label P = ( p 1 , p 2 , . . . , p m ) is R R
R
R
P P P P
P
P
P
P
P
Preprocessing
Dataset
Build a k -NN graph G of P .
2: Run a balanced graph partitioning algorithm on G into m parts.
Number the parts arbitrarily as 1 , . . .
, m .
Let π ( p ) ∈ { 1 , . . .
, m } denote the part containing p , for every p ∈ P .
3: Train a machine learning model M with training set P and labels { π ( p ) } p ∈ P .
For every x ∈ R d , let M ( x ) ∈ { 1 , . . .
, m } denote the prediction of M on x .
M ( · ) deﬁnes our m -way partition of R d .
Note that it is possible that π ( p ) (cid:54) =
M ( p ) for some p ∈ P , if M attains imperfect training
accuracy
Query
Input: query point q ∈ R d , number of bins to search b 1: Run inference on M to compute M ( q ) .
Search for
a near neighbor of q in the bin M ( q ) , i.e., among the candidates { p ∈ P
M ( p ) =
M ( q ) } .
3: If M furthermore predicts a distribution over bins, search for a near neighbor in the b top-ranked bins according to the ranking induced by the distribution (i.e., from the most likely bin to less likely ones).
Algorithm 1: space partition a distribution over the bin containing a point chosen uniformly at random among S nearest neighbors of p (including p itself).
Now, for a predicted distribution Q = ( q 1 , q 2 , . . .
, q m ) , we seek to minimize the KL divergence between P and
Q : (cid:80) mi =1 p i log p
i q i .
Intuitively, soft labels help guide the neural network with information about multiple bin ranking.
S is a hyperparameter that needs to be tuned; we study its setting in the appendix (cf.
Figure 6b).
Sparse hyperplane-induced cuts in k -NN graphs
We state and prove a theorem that shows, under certain mild assumptions, that the k -NN graph of a dataset P ⊆
R d can be partitioned by a hyperplane such that the induced cut is sparse (i.e., has few crossing edges while the sizes of two parts are similar).
The theorem is based on the framework of (Andoni et al., 2018b;c) and uses spectral techniques.
We start with some notation.
Let N k ( p ) be the set of k nearest neighbors of p in P .
The degree of p in the k -NN graph is deg( p ) =
| N k ( p ) ∪ { p (cid:48)
∈ P | p ∈ N k
( p (cid:48) ) }
| .
Let D be the distribution over the dataset P , where a point p ∈ P is sampled with probability proportional to its degree deg( p ) .
Let D close be the distribution over pairs ( p, p (cid:48) )
∈ P × P , where p ∈ P is uniformly random, and p (cid:48) is a uniformly random element of N k ( p ) .
Denote α = E ( p,p (cid:48) ) ∈D close [ (cid:107) p − p (cid:48)
(cid:107) 22 ] and β =
E x 1 ∼D ,x 2 ∼D [ (cid:107) p 1 − p 2 (cid:107) 2 2 ] .
We will proceed assuming that α (typical distance between a data point and its nearest neighbors) is noticeably smaller than β (typical distance between two independent data points).
The following theorem implies, informally speaking, that if α (cid:28) β , then there exists a hyperplane which splits the dataset into two parts of not too different size while separating only few pairs of ( p, p (cid:48) ) , where p (cid:48) is
Theorem 3.1.
There exists a hyperplane H = { x ∈ R d | (cid:104) a, x (cid:105)
=
b } such that the following holds.
Let P =
P 1 ∪ Pbe the partition of P induced by H : P 1 = { p ∈ P
| (cid:104) a, p (cid:105) ≤ b } , P 2 = { p ∈ P | (cid:104) a, p (cid:105) >
b } .
Then, one has:
Pr ( p,p (cid:48) )
∼D close [ p and p (cid:48) are separated by H ]
[Equation]
Experiments
Datasets For the experimental evaluation, we use three standard ANN benchmarks (Aum¨uller et al., 2017): SIFT (image descriptors, 1M 128-dimensional points), GloVe (word embeddings (Pennington et al., 2014), approximately 1.2M 100-dimensional points, normalized), and MNIST (images of digits, 60K 784-dimensional points).
All three datasets come with 10 000 query points, which are used for evaluation.
We include the results for SIFT and GloVe in the main text, and MNIST in Appendix A.
Evaluation metrics We mainly investigate the trade-off between the number of candidates generated for a query point, and the k -NN accuracy, deﬁned as the fraction of its k nearest neighbors that are among those candidates.
The number of candidates determines the processing time of an individual query.
Over the entire query set, we report both the average as well as the 0 .
95 -th quantile of the number of candidates.
The former measures the throughput 1 of the data structure, while the latter measures its latency.
2 We focus on parameter regimes that yield k -NN accuracy of at least 0 .
75 , in the setting k = 10 .
Additional results with broader regimes of accuracy and of k are included in the appendix.
Our methods We evaluate two variants of our method, with two different choices of the supervised learning component:
• Neural LSH
In this variant we use small neural networks.
We compare this method with k -means clustering, Iterative Quantization (ITQ)
(Gong et al., 2013), Cross-polytope LSH (Andoni et al., 2015), and Neural Catalyzer (Sablayrolles et al., 2019) composed over k -means clustering.
We evaluate partitions into 16 bins andbins.
We test both one-level (non-hierarchical) and two-level (hierarchical) partitions.
Queries are multi-probe.
• Regression LSH:
This variant uses logistic regression as the supervised learning component and, as a result, produces very simple partitions induced by hyperplanes.
We compare this method with PCA trees (Sproull, 1991; Kumar et al., 2008; Abdullah et al., 2014), random projection trees (Dasgupta & Sinha, 2013), and recursive bisections using 2 -means clustering.
We build trees of hierarchical bisections of depth up to 10 (thus total number of leaves up to 1024 ).
The query procedure descends a single root-to-leaf path and returns the candidates in that leaf.
4. Implementation details
Neural LSH uses a ﬁxed neural network architecture for the top-level partition, and a ﬁxed architecture for all secondlevel partitions.
Both architectures consist of several blocks, where each block is a fully-connected layer + batch normalization (Ioffe & Szegedy, 2015) + ReLU activations.
The ﬁnal block is followed by a fully-connected layer and a softmax layer.
The resulting network predicts a distribution over the bins of the partition.
The only difference between the top-level network the second-level network architecture is their number of blocks ( b ) and the size of their hidden layers ( s ).
In the top-level network we use b = 3 and s = 512 .
In the second-level networks we use b = 2 and s = 390 .
To reduce overﬁtting, we use dropout with probability 0 .
1 during training.
The networks are trained using the Adam optimizer (Kingma & Ba, 2015) for under 20 epochs on both levels.
We reduce the learning rate multiplicatively at regular intervals.
The weights are initialized with Glorot initialization (Glorot & Bengio, 2010).
To tune soft labels, we try different values of S between 1 and 120 .
We evaluate two settings for the number of bins in each level, m = 16 and m = 256 (leading to a total number of bins of the total number of bins in the two-level experiments are 16 2 = 256 and 256 2 =
65 536 , respectively).
In the two-level setting with m = 256 the bottom level of Neural LSH uses k -means instead of a neural network, to avoid overﬁtting when the number of points per bin is tiny.
The other conﬁgurations (two-levels with m = 16 and one-level with either m = 16 or m = 256 ) we use Neural LSH at all levels.
We slightly modify the KaHIP partitioner to make it more efﬁcient on the k -NN graphs.
Namely, we introduce a hard threshold of 2000 on the number of iterations for the local search part of the algorithm, which speeds up the partitioning dramatically, while barely affecting the quality of the resulting partitions.
[Figure]
[Table]
4. Comparison with multi-bin methods
Figure 4 shows the empirical comparison of Neural LSH with k -means clustering, ITQ, Cross-polytope LSH, and
Neural Catalyzer composed over k -means clustering.
It turns out that k -means is the strongest among these baselines.
The points depicted in Figure 4 are those that attain accuracy ≥ 0 .
In the appendix (Figure 10) we include the full accuracy range for all methods.
In all settings considered, Neural LSH yields consistently better partitions than k -means.
Depending on the setting, k -means requires signiﬁcantly more candidates to achieve the same accuracy:
• Up to 117% more for the average number of candidates for GloVe; • Up to 130% more for the 0 .
95 -quantiles of candidates for GloVe; • Up to 18% more for the average number of candidates for SIFT; • Up to 34% more for the 0 .
95 -quantiles of candidates for SIFT;
Figure 3 lists the largest multiplicative advantage in the number of candidates of Neural LSH compared to k -means, for accuracy values of at least 0 . 85 .
Speciﬁcally, for every conﬁguration of k -means, we compute the ratio between the number of candidates in that conﬁguration and the number of candidates of Neural LSH in its optimal conﬁguration, among those that attained at least the same accuracy as that k -means conﬁguration.
We also note that in all settings except two-level partitioning with m = 256 , 5 Neural LSH produces partitions for which the 0 .
95 -quantiles for the number of candidates are very close to the average number of candidates, which indicates very little variance between query times over different query points.
In contrast, the respective gap in the partitions produced by k -means is much larger, since unlike Neural LSH, it does not directly favor balanced partitions.
This implies that Neural LSH might be particularly suitable for latency-critical NNS applications.
Model sizes.
The largest model size learned by Neural LSH is equivalent to storing about ≈ 5700 points for SIFT, or ≈ 7100 points for GloVe.
This is considerably larger than k -means with k ≤ 256 , which stores at most 256 points.
Nonetheless, we believe the larger model size is acceptable for Neural LSH, for the following reasons.
First, in most of the NNS applications, especially for the distributed setting, the bottleneck in the high accuracy regime is the memory accesses needed to retrieve candidates and the further processing (such as distance computations, exact or approximate).
The model size is not a hindrance as long as does not exceed certain reasonable limits (e.g., it should ﬁt into a CPU cache).
Neural LSH signiﬁcantly reduces the memory access cost, while increasing the model size by an acceptable amount.
Second, we have observed that the quality of the Neural LSH partitions is not too sensitive to decreasing the sizes the hidden layers.
The model sizes we report are, for the sake of concreteness, the largest ones that still lead to improved performance.
Larger models do not increase the accuracy, and sometimes decrease it due to overﬁtting. [Figure]
4. Comparison with tree-based methods
Next we compare binary decision trees, where in each tree node a hyperplane is used to determine which of the two subtrees to descend into.
We generate hyperplanes with the following methods: Regression LSH, the Learned KD-tree of Cayton & Dasgupta (2007), the Boosted Search Forest of Li et al. (2011), cutting the dataset into two equal halves along the top PCA direction (Sproull, 1991; Kumar et al., 2008), 2 -means clustering, and random projections of the centered dataset (Dasgupta & Sinha, 2013; Keivani & Sinha, 2018).
We build trees of depth up to 10 , which correspond to hierarchical partitions with the up to 2 10 = 1024 bins.
Results for GloVe and SIFT are summarized in Figure 5 (see appendix).
For random projections, we run each conﬁguration 30 times and average the results.
For GloVe, Regression LSH signiﬁcantly outperforms 2 -means, while for SIFT, Regression LSH essentially matches 2 -means in terms of the average number of candidates, but shows a noticeable advantage in terms of the 0 .
95 -percentiles.
In both instances, Regression LSH signiﬁcantly outperforms PCA tree, and all of the above methods dramatically improve upon random projections.
Note, however, that random projections have an additional beneﬁt: in order to boost search accuracy, one can simply repeat the sampling process several times and generate an ensemble of decision trees instead of a single tree.
This allows making each individual tree relatively deep, which decreases the overall number of candidates, trading space for query time.
Other considered approaches (Regression LSH, 2 -means, PCA tree) are inherently deterministic, and boosting their accuracy requires more care: for instance, one can use partitioning into blocks as in J
´ egou et al.
(2011), or alternative approaches like Keivani & Sinha (2018).
Since we focus on individual partitions and not ensembles, we leave this issue out of the scope.
4. Additional experiments
In this section we include several additional experiments.
First, we study the effect of setting k .
We evaluate the 50 -NN accuracy of Neural LSH when the partitioning step is run on either the 10 -NN or the 50 -NN graph.[-
-]
6 We compare both algorithms to k -means with k = 50 .
Figure 6a compares these three algorithms on GloVe for 16 bins reporting average numbers of candidates.
From this plot, we can see that for k = 50 , Neural LSH convincingly outperforms k -means, and whether we use 10 -NN or 50 -NN graph matters very little.
Second, we study the effect of varying S (the soft labels parameter) for Neural LSH on GloVe for 256 bins. See Figure 6b where we report the average number of candidates.
As we can see from the plot, the setting S = 15 yields much better results compared to the vanilla case of S = 1 . However, increasing S beyond 15 brings diminishing returns on the overall accuracy. [Figure]
Conclusions and future directions
We presented a new technique for ﬁnding partitions of R d which support high-performance indexing for sublinear-time
NNS.
It proceeds in two major steps: (1) We perform a combinatorial balanced partitioning of the k -NN graph of the dataset; (2) We extend the resulting partition to the whole ambient space R d by using supervised classiﬁcation (such as logistic regression, neural networks, etc.).
Our experiments show that the new approach consistently outperforms quantization-based and tree-based partitions.
There is a number of exciting open problems we would like to highlight:
• Can we use our approach for NNS over non-Euclidean geometries, such as the edit distance (Zhang & Zhang, 2017) or the optimal transport distance (Kusner et al., 2015)?
The graph partitioning step directly carries through, but the learning step may need to be adjusted.
• Can we jointly optimize a graph partition and a classiﬁer at the same time?
By making the two components aware of each other, we expect the quality of the resulting partition of R d to improve.
A related approach has been successfully applied in Li et al.
(2011) for hyperplane tree partitions.
• Can our approach be extended to learning several high-quality partitions that complement each other?
Such an ensemble might be useful to trade query time for memory usage (Andoni et al., 2017).
• Can we use machine learning techniques to improve graph-based indexing techniques (Malkov & Yashunin, 2018) for NNS?
(This is in contrast to partition-based indexing, as done in this work).
• Our framework is an example of combinatorial tools aiding “continuous” learning techniques.
A more open-ended question is whether other problems can beneﬁt from such symbiosis.
A Results for MNIST
We include experimental results for the MNIST dataset, where all the experiments are performed exactly in the same way as for SIFT and GloVe.
Consistent with the trend we observed for SIFT and GloVe, Neural LSH consistently outperforms k -means (see Figure 7) both in terms of average number of candidates and especially in terms of the 0 .
95 -th quantiles.
We also compare Regression LSH with recursive 2 -means, as well as PCA tree and random projections (see
Figure 8) where Regression LSH consistently outperforms the other methods.
[Figure]
[Figure]
B Effect of Neural Catalyzer on Space Partitions
In this section we compare vanilla k -means with k -means run after applying a Neural Catalyzer map (Sablayrolles et al., 2019).
The goal is to check whether the Neural Catalyzer – which is designed to boost up the performance of sketching methods for NNS by adjusting the input geometry – could also improve the quality of space partitions for NNS.
See Figure 9 for the comparison on GloVe and SIFT with 16 bins.
On both datasets (especially SIFT), Neural Catalyzer in fact degrades the quality of the partitions.
We observed a similar trend for other numbers of bins than the setting reported here.
These ﬁndings support our observation that while both indexing and sketching for NNS can beneﬁt from learning-based enhancements, they are fundamentally different approaches and require different specialized techniques.
[Figure]
C Proof of Theorem 3.
Proof.
Consider an undirected graph G =
( V, E ) , where the set of vertices V is P , and the (multi-)set of edges contains an edge ( p, p (cid:48) )
for every p (cid:48) ∈ N k ( p ) .
The graph contains n vertices and kn edges, and some of the edges might be double (if p (cid:48) ∈ N k ( p ) and p ∈ N k
( p (cid:48) )
at the same time).
Let A G be the symmetric adjacency matrix of G normalized by 2 kn (so that the sum of all the entries equals to 1 , thus giving a probability distribution over P × P , which can be seen to be equal to D close ).
The rows and columns of A G can naturally be indexed by the points of P .
Denote ρ G ( p ) =
(cid:80) p (cid:48)
( A G ) p,p (cid:48) .
It is immediate to check that ρ G yields a distribution over P , which can be seen to be equal to
D .
Denote D G = diag( ρ G ) .
Denote L G =
D G − A G the Laplacian of A G .
Due to the equivalence of ρ G and D and
A G and D close , we have:
[Equation]
By considering all possible coordinate projections and using additivity of (cid:107) · (cid:107) 22 over coordinates, we conclude that there exists a coordinate i ∗ ∈
[ d ] such that:
[Equation]
Deﬁne a vector y ∈ R P by y p = p
i ∗ .
We now apply the following standard fact from spectral graph theory:
If A is the weighted adjacency matrix of a graph, and L is its Laplacian matrix , then x t Lx = (cid:80) ni,j =1 A ij ( x i − x j ) 2 for all x ∈ R n .
Thus the numerator of (3) becomes y t L G y .
For the denominator, consider the graph H on P in which every pair p, p (cid:48) is connected by an edge of weight ρ G ( p ) ρ G ( p (cid:48) ) .
• Its weighted adjacency matrix
A H is given by ( A H ) p,p (cid:48)
=
ρ G ( p ) ρ G ( p (cid:48) ) for p (cid:54) = p (cid:48) , and with zeros on the diagonal.
Thus A H =
ρ G ρ tG − D 2 G (recall that D G = diag( ρ G ) ).
• The degree of each node p in H equals ρ G ( p ) (cid:80) p (cid:48)
∈ P \{ p } ρ G ( p (cid:48) )
=
ρ G ( p ) − ( ρ G ( p )) 2 (recall that (cid:80) p (cid:48)
∈ P ρ G ( p (cid:48) )
= 1 ).
Therefore the diagonal degree matrix of H is D H
=
D G − D 2 G .
[Figure]
Together, the Lapacian of H is L H =
D H − A H =
D G − ρ G ρ tG .
Therefore the denominator of (3) becomes y t ( D G − ρ G ρ tG ) y .
Overall, we have:
[Equation]
Next, we deﬁne (cid:101) y
= y − c · 1 , where 1 is the all- 1 ’s vector, and c is the scalar c = ( y t ρ G ) / ( 1 t ρ G ) .
This scalar is chosen to render (cid:101) y ⊥
ρ G .
Furthermore, since 1 is in the kernel of every Laplacian matrix, we have L G (cid:101) y
=
L G y and
L H (cid:101) y =
L H y .
Together, we get
[Equation]
Now by the Cheeger’s inequality (Chung, 1996), we conclude that there exists a threshold y
∈
R such that:
[Equation]
One can trace back all the deﬁnitions and observe that the set { p ∈ P : (cid:101) y p ≤ y 0 } is induced by an (axis-aligned) hyperplane, and the left-hand side of (4) is nothing else but the